fix: add kms_key optional arg to Pipeline.deploy() (#1962)

* fix: propagate kms_key to Pipeline.deploy()

* fix kms_key docstring

Co-authored-by: Chuyang Deng <[email protected]>
Co-authored-by: icywang86rui <[email protected]>

Author: 3 people

src/sagemaker/pipeline.py

@@ -111,6 +111,7 @@ def deploy(
         wait=True,
         update_endpoint=False,
         data_capture_config=None,
+        kms_key=None,
     ):
         """Deploy the ``Model`` to an ``Endpoint``.
 
@@ -156,6 +157,9 @@ def deploy(
             data_capture_config (sagemaker.model_monitor.DataCaptureConfig): Specifies
                 configuration related to Endpoint data capture for use with
                 Amazon SageMaker Model Monitoring. Default: None.
+            kms_key (str): The ARN, Key ID or Alias 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
@@ -192,6 +196,7 @@ def deploy(
                 initial_instance_count=initial_instance_count,
                 instance_type=instance_type,
                 tags=tags,
+                kms_key=kms_key,
                 data_capture_config_dict=data_capture_config_dict,
             )
             self.sagemaker_session.update_endpoint(
@@ -202,6 +207,7 @@ def deploy(
                 name=self.endpoint_name,
                 production_variants=[production_variant],
                 tags=tags,
+                kms_key=kms_key,
                 wait=wait,
                 data_capture_config_dict=data_capture_config_dict,
             )

tests/unit/test_pipeline_model.py

@@ -113,7 +113,8 @@ def test_deploy(tfo, time, sagemaker_session):
     model = PipelineModel(
         models=[framework_model, sparkml_model], role=ROLE, sagemaker_session=sagemaker_session
     )
-    model.deploy(instance_type=INSTANCE_TYPE, initial_instance_count=1)
+    kms_key = "pipeline-model-deploy-kms-key"
+    model.deploy(instance_type=INSTANCE_TYPE, initial_instance_count=1, kms_key=kms_key)
     sagemaker_session.endpoint_from_production_variants.assert_called_with(
         name="mi-1-2017-10-10-14-14-15",
         production_variants=[
@@ -126,6 +127,7 @@ def test_deploy(tfo, time, sagemaker_session):
             }
         ],
         tags=None,
+        kms_key=kms_key,
         wait=True,
         data_capture_config_dict=None,
     )
@@ -154,6 +156,7 @@ def test_deploy_endpoint_name(tfo, time, sagemaker_session):
             }
         ],
         tags=None,
+        kms_key=None,
         wait=True,
         data_capture_config_dict=None,
     )
@@ -183,6 +186,7 @@ def test_deploy_update_endpoint(tfo, time, sagemaker_session):
         initial_instance_count=INSTANCE_COUNT,
         instance_type=INSTANCE_TYPE,
         tags=None,
+        kms_key=None,
         data_capture_config_dict=None,
     )
     config_name = sagemaker_session.create_endpoint_config(
@@ -275,6 +279,7 @@ def test_deploy_tags(tfo, time, sagemaker_session):
         ],
         tags=tags,
         wait=True,
+        kms_key=None,
         data_capture_config_dict=None,
     )