diff --git a/src/sagemaker/automl/automl.py b/src/sagemaker/automl/automl.py index f6782ce72a..eddd8de698 100644 --- a/src/sagemaker/automl/automl.py +++ b/src/sagemaker/automl/automl.py @@ -420,6 +420,7 @@ def deploy( serializer=serializer, deserializer=deserializer, endpoint_name=endpoint_name, + kms_key=model_kms_key, tags=tags, wait=wait, ) diff --git a/tests/unit/sagemaker/automl/test_auto_ml.py b/tests/unit/sagemaker/automl/test_auto_ml.py index 8edcdd048c..1b25fef693 100644 --- a/tests/unit/sagemaker/automl/test_auto_ml.py +++ b/tests/unit/sagemaker/automl/test_auto_ml.py @@ -541,6 +541,7 @@ def test_deploy(sagemaker_session, candidate_mock): initial_instance_count=INSTANCE_COUNT, instance_type=INSTANCE_TYPE, sagemaker_session=sagemaker_session, + model_kms_key=OUTPUT_KMS_KEY, ) auto_ml.create_model.assert_called_once() mock_pipeline.deploy.assert_called_once() @@ -594,6 +595,7 @@ def test_deploy_optional_args(candidate_estimator, sagemaker_session, candidate_ serializer=None, deserializer=None, endpoint_name=JOB_NAME, + kms_key=OUTPUT_KMS_KEY, tags=TAGS, wait=False, )