Skip to content
Merged
Show file tree
Hide file tree
Changes from 5 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion src/sagemaker/jumpstart/hub/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -106,7 +106,7 @@ def construct_hub_model_reference_arn_from_inputs(
info = get_info_from_hub_resource_arn(hub_arn)
arn = (
f"arn:{info.partition}:sagemaker:{info.region}:{info.account_id}:hub-content/"
f"{info.hub_name}/{HubContentType.MODEL_REFERENCE}/{model_name}/{version}"
f"{info.hub_name}/{HubContentType.MODEL_REFERENCE.value}/{model_name}/{version}"
)

return arn
Expand Down
14 changes: 14 additions & 0 deletions tests/unit/sagemaker/jumpstart/hub/test_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -95,6 +95,20 @@ def test_construct_hub_model_arn_from_inputs():
== "arn:aws:sagemaker:us-west-2:123456789123:hub-content/my-mock-hub/Model/pytorch-ic-imagenet-v2/*"
)

def test_construct_hub_model_reference_arn_from_inputs():
model_name, version = "pytorch-ic-imagenet-v2", "1.0.2"
hub_arn = "arn:aws:sagemaker:us-west-2:123456789123:hub/my-mock-hub"

assert (
utils.construct_hub_model_arn_reference_from_inputs(hub_arn, model_name, version)
== "arn:aws:sagemaker:us-west-2:123456789123:hub-content/my-mock-hub/ModelReference/pytorch-ic-imagenet-v2/1.0.2"
)

version = "*"
assert (
utils.construct_hub_model_arn_reference_from_inputs(hub_arn, model_name, version)
== "arn:aws:sagemaker:us-west-2:123456789123:hub-content/my-mock-hub/ModelReference/pytorch-ic-imagenet-v2/*"
)

def test_generate_hub_arn_for_init_kwargs():
hub_name = "my-hub-name"
Expand Down
Loading