diff --git a/src/sagemaker/jumpstart/utils.py b/src/sagemaker/jumpstart/utils.py index dfe3d7f1dd..d5c769efe0 100644 --- a/src/sagemaker/jumpstart/utils.py +++ b/src/sagemaker/jumpstart/utils.py @@ -1595,9 +1595,10 @@ def _add_model_access_configs_to_model_data_sources( ) acked_model_data_sources.append(mutable_model_data_source) else: - mutable_model_data_source.pop( - "HostingEulaKey" - ) # pop when model access config is not applicable + if "HostingEulaKey" in mutable_model_data_source: + mutable_model_data_source.pop( + "HostingEulaKey" + ) # pop when model access config is not applicable acked_model_data_sources.append(mutable_model_data_source) return acked_model_data_sources diff --git a/tests/unit/sagemaker/jumpstart/test_utils.py b/tests/unit/sagemaker/jumpstart/test_utils.py index 67681e2b7b..d228b4450e 100644 --- a/tests/unit/sagemaker/jumpstart/test_utils.py +++ b/tests/unit/sagemaker/jumpstart/test_utils.py @@ -2318,6 +2318,28 @@ def test_multiple_gated_additional_model_data_source_should_accept_both(self): + self.MOCK_GATED_DEPLOY_CONFIG_ADDITIONAL_MODEL_DATA_SOURCE_POST_CALL ) + def test_gated_additional_model_data_source_already_accepted_with_no_hosting_eula_key_should_pass_through( + self, + ): + mock_gated_deploy_config_additional_model_data_pre_accepted = [ + { + "ChannelName": "draft_model", + "S3DataSource": { + "CompressionType": "None", + "S3DataType": "S3Prefix", + "S3Uri": "s3://jumpstart_bucket/path/to/gated/resources/", + "ModelAccessConfig": {"AcceptEula": True}, + }, + } + ] + + utils._add_model_access_configs_to_model_data_sources( + model_data_sources=mock_gated_deploy_config_additional_model_data_pre_accepted, + model_access_configs={self.MOCK_GATED_MODEL_ID: ModelAccessConfig(accept_eula=False)}, + model_id=self.MOCK_GATED_MODEL_ID, + region=JUMPSTART_DEFAULT_REGION_NAME, + ) + # Mixed Positive Cases def test_multiple_mixed_additional_model_data_source_should_pass_through_one_accept_the_other(