Fix changelog and add integ test

This also adds volume KMS key to Estimator.transformer()

Author: Andrew Packer

CHANGELOG.rst

@@ -6,12 +6,12 @@ CHANGELOG
 ========
 
 * bug-fix: Local Mode: Create output/data directory expected by SageMaker Container.
+* enhancement: Add support for volume KMS key to Transformer
 
 1.9.2
 =====
 
 * feature: add support for TensorFlow 1.9
-* enhancement: Add support for volume KMS key to Transformer
 
 1.9.1
 =====

src/sagemaker/estimator.py

@@ -328,7 +328,7 @@ def delete_endpoint(self):
 
     def transformer(self, instance_count, instance_type, strategy=None, assemble_with=None, output_path=None,
                     output_kms_key=None, accept=None, env=None, max_concurrent_transforms=None,
-                    max_payload=None, tags=None, role=None):
+                    max_payload=None, tags=None, role=None, volume_kms_key=None):
         """Return a ``Transformer`` that uses a SageMaker Model based on the training job. It reuses the
         SageMaker Session and base job name used by the Estimator.
 
@@ -350,6 +350,8 @@ def transformer(self, instance_count, instance_type, strategy=None, assemble_wit
                 the training job are used for the transform job.
             role (str): The ``ExecutionRoleArn`` IAM Role ARN for the ``Model``, which is also used during
                 transform jobs. If not specified, the role from the Estimator will be used.
+            volume_kms_key (str): Optional. KMS key ID for encrypting the volume attached to the ML
+                compute instance (default: None).
         """
         self._ensure_latest_training_job()
 
@@ -360,7 +362,7 @@ def transformer(self, instance_count, instance_type, strategy=None, assemble_wit
                            output_path=output_path, output_kms_key=output_kms_key, accept=accept,
                            max_concurrent_transforms=max_concurrent_transforms, max_payload=max_payload,
                            env=env, tags=tags, base_transform_job_name=self.base_job_name,
-                           sagemaker_session=self.sagemaker_session)
+                           volume_kms_key=volume_kms_key, sagemaker_session=self.sagemaker_session)
 
     @property
     def training_job_analytics(self):
@@ -756,7 +758,7 @@ def _update_init_params(cls, hp, tf_arguments):
 
     def transformer(self, instance_count, instance_type, strategy=None, assemble_with=None, output_path=None,
                     output_kms_key=None, accept=None, env=None, max_concurrent_transforms=None,
-                    max_payload=None, tags=None, role=None, model_server_workers=None):
+                    max_payload=None, tags=None, role=None, model_server_workers=None, volume_kms_key=None):
         """Return a ``Transformer`` that uses a SageMaker Model based on the training job. It reuses the
         SageMaker Session and base job name used by the Estimator.
 
@@ -780,6 +782,8 @@ def transformer(self, instance_count, instance_type, strategy=None, assemble_wit
                 transform jobs. If not specified, the role from the Estimator will be used.
             model_server_workers (int): Optional. The number of worker processes used by the inference server.
                 If None, server will use one worker per vCPU.
+            volume_kms_key (str): Optional. KMS key ID for encrypting the volume attached to the ML
+                compute instance (default: None).
         """
         self._ensure_latest_training_job()
         role = role or self.role
@@ -799,7 +803,7 @@ def transformer(self, instance_count, instance_type, strategy=None, assemble_wit
                            output_path=output_path, output_kms_key=output_kms_key, accept=accept,
                            max_concurrent_transforms=max_concurrent_transforms, max_payload=max_payload,
                            env=transform_env, tags=tags, base_transform_job_name=self.base_job_name,
-                           sagemaker_session=self.sagemaker_session)
+                           volume_kms_key=volume_kms_key, sagemaker_session=self.sagemaker_session)
 
 
 def _s3_uri_prefix(channel_name, s3_data):

tests/integ/kms_utils.py

@@ -0,0 +1,89 @@
+# Copyright 2017-2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+from __future__ import absolute_import
+
+KEY_ALIAS = "SageMakerKmsKey"
+KEY_POLICY = '''
+{{
+  "Version": "2012-10-17",
+  "Id": "sagemaker-kms-integ-test-policy",
+  "Statement": [
+    {{
+      "Sid": "Enable IAM User Permissions",
+      "Effect": "Allow",
+      "Principal": {{
+        "AWS": "arn:aws:iam::{account_id}:root"
+      }},
+      "Action": "kms:*",
+      "Resource": "*"
+    }},
+    {{
+      "Sid": "Allow use of the key",
+      "Effect": "Allow",
+      "Principal": {{
+        "AWS": "arn:aws:iam::{account_id}:role/SageMakerRole"
+      }},
+      "Action": [
+        "kms:Encrypt",
+        "kms:Decrypt",
+        "kms:ReEncrypt*",
+        "kms:GenerateDataKey*",
+        "kms:DescribeKey"
+      ],
+      "Resource": "*"
+    }},
+    {{
+      "Sid": "Allow attachment of persistent resources",
+      "Effect": "Allow",
+      "Principal": {{
+        "AWS": "arn:aws:iam::{account_id}:role/SageMakerRole"
+      }},
+      "Action": [
+        "kms:CreateGrant",
+        "kms:ListGrants",
+        "kms:RevokeGrant"
+      ],
+      "Resource": "*",
+      "Condition": {{
+        "Bool": {{
+          "kms:GrantIsForAWSResource": "true"
+        }}
+      }}
+    }}
+  ]
+}}
+'''
+
+
+def _get_kms_key_arn(kms_client, alias):
+    try:
+        response = kms_client.describe_key(KeyId='alias/' + alias)
+        return response['KeyMetadata']['Arn']
+    except kms_client.exceptions.NotFoundException:
+        return None
+
+def _create_kms_key(kms_client, account_id):
+    response = kms_client.create_key(
+        Policy=KEY_POLICY.format(account_id=account_id),
+        Description='KMS key for SageMaker Python SDK integ tests',
+    )
+    key_arn = response['KeyMetadata']['Arn']
+    response = kms_client.create_alias(AliasName='alias/' + KEY_ALIAS, TargetKeyId=key_arn)
+    return key_arn
+
+def get_or_create_kms_key(kms_client, account_id):
+    kms_key_arn = _get_kms_key_arn(kms_client, KEY_ALIAS)
+    if kms_key_arn is not None:
+        return kms_key_arn
+    else:
+        return _create_kms_key(kms_client, account_id)

tests/integ/test_transformer.py

@@ -23,6 +23,7 @@
 from sagemaker.mxnet import MXNet
 from sagemaker.transformer import Transformer
 from tests.integ import DATA_DIR, TRAINING_DEFAULT_TIMEOUT_MINUTES
+from tests.integ.kms_utils import get_or_create_kms_key
 from tests.integ.timeout import timeout
 
 
@@ -47,9 +48,16 @@ def test_transform_mxnet(sagemaker_session):
     transform_input = mx.sagemaker_session.upload_data(path=transform_input_path,
                                                        key_prefix=transform_input_key_prefix)
 
-    transformer = _create_transformer_and_transform_job(mx, transform_input)
-    transformer.wait()
+    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)
 
+    transformer = _create_transformer_and_transform_job(mx, transform_input, kms_key_arn)
+    transformer.wait()
+    job_desc = transformer.sagemaker_session.sagemaker_client.describe_transform_job(
+        TransformJobName=transformer.latest_transform_job.name)
+    assert kms_key_arn == job_desc['TransformResources']['VolumeKmsKeyId']
 
 @pytest.mark.continuous_testing
 def test_attach_transform_kmeans(sagemaker_session):
@@ -90,7 +98,7 @@ def test_attach_transform_kmeans(sagemaker_session):
     attached_transformer.wait()
 
 
-def _create_transformer_and_transform_job(estimator, transform_input):
-    transformer = estimator.transformer(1, 'ml.m4.xlarge')
+def _create_transformer_and_transform_job(estimator, transform_input, volume_kms_key=None):
+    transformer = estimator.transformer(1, 'ml.m4.xlarge', volume_kms_key=volume_kms_key)
     transformer.transform(transform_input, content_type='text/csv')
     return transformer

tests/unit/test_estimator.py

@@ -488,7 +488,7 @@ def test_framework_transformer_creation_with_optional_params(name_from_image, sa
     transformer = fw.transformer(INSTANCE_COUNT, INSTANCE_TYPE, strategy=strategy, assemble_with=assemble_with,
                                  output_path=OUTPUT_PATH, output_kms_key=kms_key, accept=accept, tags=TAGS,
                                  max_concurrent_transforms=max_concurrent_transforms, max_payload=max_payload,
-                                 env=env, role=new_role, model_server_workers=1)
+                                 volume_kms_key=kms_key, env=env, role=new_role, model_server_workers=1)
 
     sagemaker_session.create_model.assert_called_with(MODEL_IMAGE, new_role, MODEL_CONTAINER_DEF)
     assert transformer.strategy == strategy
@@ -501,6 +501,7 @@ def test_framework_transformer_creation_with_optional_params(name_from_image, sa
     assert transformer.env == env
     assert transformer.base_transform_job_name == base_name
     assert transformer.tags == TAGS
+    assert transformer.volume_kms_key == kms_key
 
 
 def test_ensure_latest_training_job(sagemaker_session):