Skip to content

Commit 2ff739f

Browse files
committed
fix pop on none error in jumpstart draft model flow
1 parent 2a5edb7 commit 2ff739f

File tree

2 files changed

+26
-3
lines changed

2 files changed

+26
-3
lines changed

src/sagemaker/jumpstart/utils.py

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1595,9 +1595,10 @@ def _add_model_access_configs_to_model_data_sources(
15951595
)
15961596
acked_model_data_sources.append(mutable_model_data_source)
15971597
else:
1598-
mutable_model_data_source.pop(
1599-
"HostingEulaKey"
1600-
) # pop when model access config is not applicable
1598+
if "HostingEulaKey" in mutable_model_data_source:
1599+
mutable_model_data_source.pop(
1600+
"HostingEulaKey"
1601+
) # pop when model access config is not applicable
16011602
acked_model_data_sources.append(mutable_model_data_source)
16021603
return acked_model_data_sources
16031604

tests/unit/sagemaker/jumpstart/test_utils.py

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2318,6 +2318,28 @@ def test_multiple_gated_additional_model_data_source_should_accept_both(self):
23182318
+ self.MOCK_GATED_DEPLOY_CONFIG_ADDITIONAL_MODEL_DATA_SOURCE_POST_CALL
23192319
)
23202320

2321+
def test_gated_additional_model_data_source_already_accepted_with_no_hosting_eula_key_should_pass_through(
2322+
self,
2323+
):
2324+
mock_gated_deploy_config_additional_model_data_pre_accepted = [
2325+
{
2326+
"ChannelName": "draft_model",
2327+
"S3DataSource": {
2328+
"CompressionType": "None",
2329+
"S3DataType": "S3Prefix",
2330+
"S3Uri": "s3://jumpstart_bucket/path/to/gated/resources/",
2331+
"ModelAccessConfig": {"AcceptEula": True},
2332+
},
2333+
}
2334+
]
2335+
2336+
utils._add_model_access_configs_to_model_data_sources(
2337+
model_data_sources=mock_gated_deploy_config_additional_model_data_pre_accepted,
2338+
model_access_configs={self.MOCK_GATED_MODEL_ID: ModelAccessConfig(accept_eula=False)},
2339+
model_id=self.MOCK_GATED_MODEL_ID,
2340+
region=JUMPSTART_DEFAULT_REGION_NAME,
2341+
)
2342+
23212343
# Mixed Positive Cases
23222344

23232345
def test_multiple_mixed_additional_model_data_source_should_pass_through_one_accept_the_other(

0 commit comments

Comments
 (0)