fix: expose kms_key parameter for deploying from training and hyperparameter tuning jobs (#1044)

Author: laurenyu

src/sagemaker/estimator.py

@@ -471,6 +471,7 @@ def deploy(
         update_endpoint=False,
         wait=True,
         model_name=None,
+        kms_key=None,
         **kwargs
     ):
         """Deploy the trained model to an Amazon SageMaker endpoint and return a
@@ -510,6 +511,9 @@ def deploy(
                 For more information about tags, see
                 https://boto3.amazonaws.com/v1/documentation\
                 /api/latest/reference/services/sagemaker.html#SageMaker.Client.add_tags
+            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.
             **kwargs: Passed to invocation of ``create_model()``.
                 Implementations may customize ``create_model()`` to accept
                 ``**kwargs`` to customize model creation during deploy.
@@ -543,6 +547,7 @@ def deploy(
             update_endpoint=update_endpoint,
             tags=self.tags,
             wait=wait,
+            kms_key=kms_key,
         )
 
     @property

src/sagemaker/tuner.py

@@ -429,6 +429,7 @@ def deploy(
         endpoint_name=None,
         wait=True,
         model_name=None,
+        kms_key=None,
         **kwargs
     ):
         """Deploy the best trained or user specified model to an Amazon
@@ -455,6 +456,9 @@ def deploy(
                 model completes (default: True).
             model_name (str): Name to use for creating an Amazon SageMaker
                 model. If not specified, the name of the training job is used.
+            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.
             **kwargs: Other arguments needed for deployment. Please refer to the
                 ``create_model()`` method of the associated estimator to see
                 what other arguments are needed.
@@ -475,6 +479,7 @@ def deploy(
             endpoint_name=endpoint_name,
             wait=wait,
             model_name=model_name,
+            kms_key=kms_key,
             **kwargs
         )
 

tests/unit/test_estimator.py

@@ -1879,6 +1879,37 @@ def test_generic_to_deploy_network_isolation(sagemaker_session):
     assert kwargs["enable_network_isolation"]
 
 
+@patch("sagemaker.estimator.Estimator.create_model")
+def test_generic_to_deploy_kms(create_model, sagemaker_session):
+    e = Estimator(
+        IMAGE_NAME,
+        ROLE,
+        INSTANCE_COUNT,
+        INSTANCE_TYPE,
+        output_path=OUTPUT_PATH,
+        sagemaker_session=sagemaker_session,
+    )
+    e.fit()
+
+    model = MagicMock()
+    create_model.return_value = model
+
+    endpoint_name = "foo"
+    kms_key = "key"
+    e.deploy(INSTANCE_COUNT, INSTANCE_TYPE, endpoint_name=endpoint_name, kms_key=kms_key)
+
+    model.deploy.assert_called_with(
+        instance_type=INSTANCE_TYPE,
+        initial_instance_count=INSTANCE_COUNT,
+        accelerator_type=None,
+        endpoint_name=endpoint_name,
+        update_endpoint=False,
+        tags=None,
+        wait=True,
+        kms_key=kms_key,
+    )
+
+
 def test_generic_training_job_analytics(sagemaker_session):
     sagemaker_session.sagemaker_client.describe_training_job = Mock(
         name="describe_training_job",

tests/unit/test_tuner.py

@@ -1,4 +1,4 @@
-# Copyright 2017 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# Copyright 2017-2019 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
@@ -16,7 +16,7 @@
 import os
 
 import pytest
-from mock import Mock
+from mock import Mock, patch
 
 from sagemaker import RealTimePredictor
 from sagemaker.amazon.amazon_estimator import RecordSet
@@ -656,6 +656,50 @@ def test_deploy_default(tuner):
     assert predictor.sagemaker_session == tuner.estimator.sagemaker_session
 
 
+@patch("sagemaker.estimator.Estimator.attach")
+@patch("sagemaker.tuner.HyperparameterTuner.best_training_job")
+def test_deploy_optional_params(best_training_job, estimator_attach, tuner):
+    tuner.fit()
+
+    estimator = Mock()
+    estimator_attach.return_value = estimator
+
+    training_job = "best-job-ever"
+    best_training_job.return_value = training_job
+
+    accelerator = "ml.eia1.medium"
+    endpoint_name = "foo"
+    model_name = "bar"
+    kms_key = "key"
+    kwargs = {"some_arg": "some_value"}
+
+    tuner.deploy(
+        TRAIN_INSTANCE_COUNT,
+        TRAIN_INSTANCE_TYPE,
+        accelerator_type=accelerator,
+        endpoint_name=endpoint_name,
+        wait=False,
+        model_name=model_name,
+        kms_key=kms_key,
+        **kwargs
+    )
+
+    estimator_attach.assert_called_with(
+        training_job, sagemaker_session=tuner.estimator.sagemaker_session
+    )
+
+    estimator.deploy.assert_called_with(
+        TRAIN_INSTANCE_COUNT,
+        TRAIN_INSTANCE_TYPE,
+        accelerator_type=accelerator,
+        endpoint_name=endpoint_name,
+        wait=False,
+        model_name=model_name,
+        kms_key=kms_key,
+        **kwargs
+    )
+
+
 def test_wait(tuner):
     tuner.latest_tuning_job = _TuningJob(tuner.estimator.sagemaker_session, JOB_NAME)
     tuner.estimator.sagemaker_session.wait_for_tuning_job = Mock(name="wait_for_tuning_job")