Add volume KMS key to transformer

CHANGELOG.rst

@@ -11,6 +11,7 @@ CHANGELOG
 =====
 
 * feature: add support for TensorFlow 1.9
+* enhancement: Add support for volume KMS key to Transformer
 
 1.9.1
 =====

src/sagemaker/transformer.py

@@ -23,7 +23,7 @@ class Transformer(object):
 
     def __init__(self, model_name, instance_count, instance_type, strategy=None, assemble_with=None, output_path=None,
                  output_kms_key=None, accept=None, max_concurrent_transforms=None, max_payload=None, tags=None,
-                 env=None, base_transform_job_name=None, sagemaker_session=None):
+                 env=None, base_transform_job_name=None, sagemaker_session=None, volume_kms_key=None):
         """Initialize a ``Transformer``.
 
         Args:
@@ -50,6 +50,8 @@ def __init__(self, model_name, instance_count, instance_type, strategy=None, ass
             sagemaker_session (sagemaker.session.Session): Session object which manages interactions with
                 Amazon SageMaker APIs and any other AWS services needed. If not specified, the estimator creates one
                 using the default AWS configuration chain.
+            volume_kms_key (str): Optional. KMS key ID for encrypting the volume attached to the ML
+                compute instance (default: None).
         """
         self.model_name = model_name
         self.strategy = strategy
@@ -62,6 +64,7 @@ def __init__(self, model_name, instance_count, instance_type, strategy=None, ass
 
         self.instance_count = instance_count
         self.instance_type = instance_type
+        self.volume_kms_key = volume_kms_key
 
         self.max_concurrent_transforms = max_concurrent_transforms
         self.max_payload = max_payload
@@ -159,6 +162,7 @@ def _prepare_init_params_from_job_description(cls, job_details):
         init_params['model_name'] = job_details['ModelName']
         init_params['instance_count'] = job_details['TransformResources']['InstanceCount']
         init_params['instance_type'] = job_details['TransformResources']['InstanceType']
+        init_params['volume_kms_key'] = job_details['TransformResources'].get('VolumeKmsKeyId')
         init_params['strategy'] = job_details.get('BatchStrategy')
         init_params['assemble_with'] = job_details['TransformOutput'].get('AssembleWith')
         init_params['output_path'] = job_details['TransformOutput']['S3OutputPath']
@@ -200,7 +204,8 @@ def _load_config(data, data_type, content_type, compression_type, split_type, tr
         output_config = _TransformJob._prepare_output_config(transformer.output_path, transformer.output_kms_key,
                                                              transformer.assemble_with, transformer.accept)
 
-        resource_config = _TransformJob._prepare_resource_config(transformer.instance_count, transformer.instance_type)
+        resource_config = _TransformJob._prepare_resource_config(transformer.instance_count, transformer.instance_type,
+                                                                 transformer.volume_kms_key)
 
         return {'input_config': input_config,
                 'output_config': output_config,
@@ -241,5 +246,5 @@ def _prepare_output_config(s3_path, kms_key_id, assemble_with, accept):
         return config
 
     @staticmethod
-    def _prepare_resource_config(instance_count, instance_type):
-        return {'InstanceCount': instance_count, 'InstanceType': instance_type}
+    def _prepare_resource_config(instance_count, instance_type, volume_kms_key):
+        return {'InstanceCount': instance_count, 'InstanceType': instance_type, 'VolumeKmsKeyId': volume_kms_key}

tests/unit/test_transformer.py

@@ -23,6 +23,7 @@
 
 INSTANCE_COUNT = 1
 INSTANCE_TYPE = 'ml.m4.xlarge'
+KMS_KEY_ID = 'kms-key-id'
 
 S3_DATA_TYPE = 'S3Prefix'
 S3_BUCKET = 'bucket'
@@ -48,7 +49,8 @@ def sagemaker_session():
 @pytest.fixture()
 def transformer(sagemaker_session):
     return Transformer(MODEL_NAME, INSTANCE_COUNT, INSTANCE_TYPE,
-                       output_path=OUTPUT_PATH, sagemaker_session=sagemaker_session)
+                       output_path=OUTPUT_PATH, sagemaker_session=sagemaker_session,
+                       volume_kms_key=KMS_KEY_ID)
 
 
 @patch('sagemaker.transformer._TransformJob.start_new')
@@ -178,7 +180,8 @@ def test_prepare_init_params_from_job_description_all_keys(transformer):
         'ModelName': MODEL_NAME,
         'TransformResources': {
             'InstanceCount': INSTANCE_COUNT,
-            'InstanceType': INSTANCE_TYPE
+            'InstanceType': INSTANCE_TYPE,
+            'VolumeKmsKeyId': KMS_KEY_ID
         },
         'BatchStrategy': None,
         'TransformOutput': {
@@ -197,6 +200,7 @@ def test_prepare_init_params_from_job_description_all_keys(transformer):
     assert init_params['model_name'] == MODEL_NAME
     assert init_params['instance_count'] == INSTANCE_COUNT
     assert init_params['instance_type'] == INSTANCE_TYPE
+    assert init_params['volume_kms_key'] == KMS_KEY_ID
 
 
 # _TransformJob tests
@@ -227,6 +231,7 @@ def test_load_config(transformer):
         'resource_config': {
             'InstanceCount': INSTANCE_COUNT,
             'InstanceType': INSTANCE_TYPE,
+            'VolumeKmsKeyId': KMS_KEY_ID,
         },
     }
 
@@ -292,8 +297,8 @@ def test_prepare_output_config_with_optional_params():
 
 
 def test_prepare_resource_config():
-    config = _TransformJob._prepare_resource_config(INSTANCE_COUNT, INSTANCE_TYPE)
-    assert config == {'InstanceCount': INSTANCE_COUNT, 'InstanceType': INSTANCE_TYPE}
+    config = _TransformJob._prepare_resource_config(INSTANCE_COUNT, INSTANCE_TYPE, KMS_KEY_ID)
+    assert config == {'InstanceCount': INSTANCE_COUNT, 'InstanceType': INSTANCE_TYPE, 'VolumeKmsKeyId': KMS_KEY_ID}
 
 
 def test_transform_job_wait(sagemaker_session):