Skip to content

Commit da77745

Browse files
authored
fix: propagate KMS key to model.deploy (#2603)
1 parent 0ce904c commit da77745

File tree

2 files changed

+3
-0
lines changed

2 files changed

+3
-0
lines changed

src/sagemaker/automl/automl.py

+1
Original file line numberDiff line numberDiff line change
@@ -420,6 +420,7 @@ def deploy(
420420
serializer=serializer,
421421
deserializer=deserializer,
422422
endpoint_name=endpoint_name,
423+
kms_key=model_kms_key,
423424
tags=tags,
424425
wait=wait,
425426
)

tests/unit/sagemaker/automl/test_auto_ml.py

+2
Original file line numberDiff line numberDiff line change
@@ -541,6 +541,7 @@ def test_deploy(sagemaker_session, candidate_mock):
541541
initial_instance_count=INSTANCE_COUNT,
542542
instance_type=INSTANCE_TYPE,
543543
sagemaker_session=sagemaker_session,
544+
model_kms_key=OUTPUT_KMS_KEY,
544545
)
545546
auto_ml.create_model.assert_called_once()
546547
mock_pipeline.deploy.assert_called_once()
@@ -594,6 +595,7 @@ def test_deploy_optional_args(candidate_estimator, sagemaker_session, candidate_
594595
serializer=None,
595596
deserializer=None,
596597
endpoint_name=JOB_NAME,
598+
kms_key=OUTPUT_KMS_KEY,
597599
tags=TAGS,
598600
wait=False,
599601
)

0 commit comments

Comments
 (0)