Skip to content

Commit cfc5dea

Browse files
committed
fix(train): fall back to SageMakerPublicHub when model not found in private hub
When _get_hub_content_metadata fails to find a model in the configured private hub (e.g. sdktest), retry with SageMakerPublicHub. This handles models like meta-textgeneration-llama-3-2-1b-instruct that only exist in the public hub.
1 parent 8e060f6 commit cfc5dea

1 file changed

Lines changed: 24 additions & 8 deletions

File tree

sagemaker-train/src/sagemaker/train/common_utils/recipe_utils.py

Lines changed: 24 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -63,14 +63,30 @@ def _get_hub_content_metadata(
6363
... )
6464
>>> print(metadata['HubContentName'])
6565
"""
66-
hub_content = HubContent.get(
67-
hub_name=hub_name,
68-
hub_content_type=hub_content_type,
69-
hub_content_name=hub_content_name,
70-
region=region,
71-
session=session
72-
)
73-
66+
try:
67+
hub_content = HubContent.get(
68+
hub_name=hub_name,
69+
hub_content_type=hub_content_type,
70+
hub_content_name=hub_content_name,
71+
region=region,
72+
session=session
73+
)
74+
except Exception:
75+
if hub_name != "SageMakerPublicHub":
76+
logger.info(
77+
f"Hub content '{hub_content_name}' not found in '{hub_name}', "
78+
f"falling back to SageMakerPublicHub"
79+
)
80+
hub_content = HubContent.get(
81+
hub_name="SageMakerPublicHub",
82+
hub_content_type=hub_content_type,
83+
hub_content_name=hub_content_name,
84+
region=region,
85+
session=session
86+
)
87+
else:
88+
raise
89+
7490
# Convert to dict for easier access
7591
hub_content_dict = hub_content.__dict__
7692

0 commit comments

Comments
 (0)