feat: Curated hub improvements

src/sagemaker/jumpstart/accessors.py

@@ -14,6 +14,7 @@
 """This module contains accessors related to SageMaker JumpStart."""
 from __future__ import absolute_import
 import functools
+import logging
 from typing import Any, Dict, List, Optional
 import boto3
 
@@ -289,15 +290,6 @@ def get_model_specs(
 
         if hub_arn:
             try:
-                hub_model_arn = construct_hub_model_arn_from_inputs(
-                    hub_arn=hub_arn, model_name=model_id, version=version
-                )
-                model_specs = JumpStartModelsAccessor._cache.get_hub_model(
-                    hub_model_arn=hub_model_arn
-                )
-                model_specs.set_hub_content_type(HubContentType.MODEL)
-                return model_specs
-            except:  # noqa: E722
                 hub_model_arn = construct_hub_model_reference_arn_from_inputs(
                     hub_arn=hub_arn, model_name=model_id, version=version
                 )
@@ -307,6 +299,21 @@ def get_model_specs(
                 model_specs.set_hub_content_type(HubContentType.MODEL_REFERENCE)
                 return model_specs
 
+            except Exception as ex:
+                logging.info(
+                    "Recieved exeption while calling APIs for ContentType Model, \
+                        retrying with ContentType ModelReference: "
+                    + str(ex)
+                )
+                hub_model_arn = construct_hub_model_arn_from_inputs(
+                    hub_arn=hub_arn, model_name=model_id, version=version
+                )
+                model_specs = JumpStartModelsAccessor._cache.get_hub_model(
+                    hub_model_arn=hub_model_arn
+                )
+                model_specs.set_hub_content_type(HubContentType.MODEL)
+                return model_specs
+
         return JumpStartModelsAccessor._cache.get_specs(  # type: ignore
             model_id=model_id, version_str=version, model_type=model_type
         )

src/sagemaker/jumpstart/factory/estimator.py

@@ -29,6 +29,10 @@
     _retrieve_model_package_model_artifact_s3_uri,
 )
 from sagemaker.jumpstart.artifacts.resource_names import _retrieve_resource_name_base
+from sagemaker.jumpstart.hub.utils import (
+    construct_hub_model_arn_from_inputs,
+    construct_hub_model_reference_arn_from_inputs,
+)
 from sagemaker.session import Session
 from sagemaker.async_inference.async_inference_config import AsyncInferenceConfig
 from sagemaker.base_deserializers import BaseDeserializer
@@ -52,6 +56,7 @@
 from sagemaker.jumpstart.enums import JumpStartScriptScope, JumpStartModelType
 from sagemaker.jumpstart.factory import model
 from sagemaker.jumpstart.types import (
+    HubContentType,
     JumpStartEstimatorDeployKwargs,
     JumpStartEstimatorFitKwargs,
     JumpStartEstimatorInitKwargs,
@@ -201,6 +206,11 @@ def get_init_kwargs(
     estimator_init_kwargs = _add_region_to_kwargs(estimator_init_kwargs)
     estimator_init_kwargs = _add_instance_type_and_count_to_kwargs(estimator_init_kwargs)
     estimator_init_kwargs = _add_image_uri_to_kwargs(estimator_init_kwargs)
+    if hub_arn:
+        estimator_init_kwargs = _add_model_reference_arn_to_kwargs(kwargs=estimator_init_kwargs)
+    else:
+        estimator_init_kwargs.model_reference_arn = None
+        estimator_init_kwargs.hub_content_type = None
     estimator_init_kwargs = _add_model_uri_to_kwargs(estimator_init_kwargs)
     estimator_init_kwargs = _add_source_dir_to_kwargs(estimator_init_kwargs)
     estimator_init_kwargs = _add_entry_point_to_kwargs(estimator_init_kwargs)
@@ -422,7 +432,7 @@ def _add_sagemaker_session_to_kwargs(kwargs: JumpStartKwargs) -> JumpStartKwargs
     kwargs.sagemaker_session = (
         kwargs.sagemaker_session
         or get_default_jumpstart_session_with_user_agent_suffix(
-            kwargs.model_id, kwargs.model_version
+            kwargs.model_id, kwargs.model_version, kwargs.hub_arn
         )
     )
     return kwargs
@@ -511,7 +521,15 @@ def _add_tags_to_kwargs(kwargs: JumpStartEstimatorInitKwargs) -> JumpStartEstima
         )
 
     if kwargs.hub_arn:
-        kwargs.tags = add_hub_content_arn_tags(kwargs.tags, kwargs.hub_arn)
+        if kwargs.model_reference_arn:
+            hub_content_arn = construct_hub_model_reference_arn_from_inputs(
+                kwargs.hub_arn, kwargs.model_id, kwargs.model_version
+            )
+        else:
+            hub_content_arn = construct_hub_model_arn_from_inputs(
+                kwargs.hub_arn, kwargs.model_id, kwargs.model_version
+            )
+        kwargs.tags = add_hub_content_arn_tags(kwargs.tags, hub_content_arn=hub_content_arn)
 
     return kwargs
 
@@ -535,6 +553,33 @@ def _add_image_uri_to_kwargs(kwargs: JumpStartEstimatorInitKwargs) -> JumpStartE
     return kwargs
 
 
+def _add_model_reference_arn_to_kwargs(
+    kwargs: JumpStartEstimatorInitKwargs,
+) -> JumpStartEstimatorInitKwargs:
+    """Sets Model Reference ARN if the hub content type is Model Reference, returns full kwargs."""
+
+    hub_content_type = verify_model_region_and_return_specs(
+        model_id=kwargs.model_id,
+        version=kwargs.model_version,
+        hub_arn=kwargs.hub_arn,
+        scope=JumpStartScriptScope.TRAINING,
+        region=kwargs.region,
+        tolerate_vulnerable_model=kwargs.tolerate_vulnerable_model,
+        tolerate_deprecated_model=kwargs.tolerate_deprecated_model,
+        sagemaker_session=kwargs.sagemaker_session,
+        model_type=kwargs.model_type,
+    ).hub_content_type
+    kwargs.hub_content_type = hub_content_type if kwargs.hub_arn else None
+
+    if hub_content_type == HubContentType.MODEL_REFERENCE:
+        kwargs.model_reference_arn = construct_hub_model_reference_arn_from_inputs(
+            hub_arn=kwargs.hub_arn, model_name=kwargs.model_id, version=kwargs.model_version
+        )
+    else:
+        kwargs.model_reference_arn = None
+    return kwargs
+
+
 def _add_model_uri_to_kwargs(kwargs: JumpStartEstimatorInitKwargs) -> JumpStartEstimatorInitKwargs:
     """Sets model uri in kwargs based on default or override, returns full kwargs."""
 

src/sagemaker/jumpstart/factory/model.py

@@ -34,7 +34,10 @@
     JUMPSTART_LOGGER,
 )
 from sagemaker.model_card.model_card import ModelCard, ModelPackageModelCard
-from sagemaker.jumpstart.hub.utils import construct_hub_model_reference_arn_from_inputs
+from sagemaker.jumpstart.hub.utils import (
+    construct_hub_model_arn_from_inputs,
+    construct_hub_model_reference_arn_from_inputs,
+)
 from sagemaker.model_metrics import ModelMetrics
 from sagemaker.metadata_properties import MetadataProperties
 from sagemaker.drift_check_baselines import DriftCheckBaselines
@@ -149,12 +152,14 @@ def _add_sagemaker_session_to_kwargs(
     kwargs: Union[JumpStartModelInitKwargs, JumpStartModelDeployKwargs]
 ) -> JumpStartModelInitKwargs:
     """Sets session in kwargs based on default or override, returns full kwargs."""
+
     kwargs.sagemaker_session = (
         kwargs.sagemaker_session
         or get_default_jumpstart_session_with_user_agent_suffix(
-            kwargs.model_id, kwargs.model_version
+            kwargs.model_id, kwargs.model_version, kwargs.hub_arn
         )
     )
+
     return kwargs
 
 
@@ -265,6 +270,7 @@ def _add_model_reference_arn_to_kwargs(
     kwargs: JumpStartModelInitKwargs,
 ) -> JumpStartModelInitKwargs:
     """Sets Model Reference ARN if the hub content type is Model Reference, returns full kwargs."""
+
     hub_content_type = verify_model_region_and_return_specs(
         model_id=kwargs.model_id,
         version=kwargs.model_version,
@@ -550,7 +556,15 @@ def _add_tags_to_kwargs(kwargs: JumpStartModelDeployKwargs) -> Dict[str, Any]:
         )
 
     if kwargs.hub_arn:
-        kwargs.tags = add_hub_content_arn_tags(kwargs.tags, kwargs.hub_arn)
+        if kwargs.model_reference_arn:
+            hub_content_arn = construct_hub_model_reference_arn_from_inputs(
+                kwargs.hub_arn, kwargs.model_id, kwargs.model_version
+            )
+        else:
+            hub_content_arn = construct_hub_model_arn_from_inputs(
+                kwargs.hub_arn, kwargs.model_id, kwargs.model_version
+            )
+        kwargs.tags = add_hub_content_arn_tags(kwargs.tags, hub_content_arn=hub_content_arn)
 
     return kwargs
 

src/sagemaker/jumpstart/hub/hub.py

@@ -23,7 +23,6 @@
 from sagemaker.session import Session
 
 from sagemaker.jumpstart.constants import (
-    DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
     JUMPSTART_LOGGER,
 )
 from sagemaker.jumpstart.types import (
@@ -68,7 +67,9 @@ def __init__(
         self,
         hub_name: str,
         bucket_name: Optional[str] = None,
-        sagemaker_session: Optional[Session] = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+        sagemaker_session: Optional[
+            Session
+        ] = utils.get_default_jumpstart_session_with_user_agent_suffix(is_hub_content=True),
     ) -> None:
         """Instantiates a SageMaker ``Hub``.
 
@@ -288,7 +289,10 @@ def describe_model(
             )
 
         except Exception as ex:
-            logging.info("Recieved expection while calling APIs for ContentType Model: " + str(ex))
+            logging.info(
+                "Recieved exeption while calling APIs for ContentType Model, retrying with ContentType ModelReference: "
+                + str(ex)
+            )
             model_version = get_hub_model_version(
                 hub_model_name=model_name,
                 hub_model_type=HubContentType.MODEL_REFERENCE.value,

src/sagemaker/jumpstart/hub/utils.py

@@ -193,7 +193,11 @@ def get_hub_model_version(
     hub_model_version: Optional[str] = None,
     sagemaker_session: Session = constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
 ) -> str:
-    """Returns available Jumpstart hub model version"""
+    """Returns available Jumpstart hub model version
+
+    Raises:
+        ResourceNotFound: If the specified model is not found in the hub.
+    """
 
     try:
         hub_content_summaries = sagemaker_session.list_hub_content_versions(

src/sagemaker/jumpstart/model.py

@@ -429,6 +429,7 @@ def attach(
         cls,
         endpoint_name: str,
         inference_component_name: Optional[str] = None,
+        hub_name: Optional[str] = None,
         model_id: Optional[str] = None,
         model_version: Optional[str] = None,
         sagemaker_session=DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
@@ -457,6 +458,7 @@ def attach(
             model_id=model_id,
             model_version=model_version,
             sagemaker_session=sagemaker_session,
+            hub_name=hub_name,
         )
         model.endpoint_name = endpoint_name
         model.inference_component_name = inference_component_name

src/sagemaker/jumpstart/types.py

@@ -1412,6 +1412,7 @@ def __init__(self, spec: Dict[str, Any], is_hub_content: Optional[bool] = False)
 
         Args:
             spec (Dict[str, Any]): Dictionary representation of spec.
+            is_hub_content (Optional[bool]): Whether the model is from a private hub.
         """
         super().__init__(spec, is_hub_content)
         self.from_json(spec)
@@ -2025,6 +2026,8 @@ class JumpStartEstimatorInitKwargs(JumpStartKwargs):
         "enable_infra_check",
         "enable_remote_debug",
         "enable_session_tag_chaining",
+        "hub_content_type",
+        "model_reference_arn",
     ]
 
     SERIALIZATION_EXCLUSION_SET = {
@@ -2035,6 +2038,7 @@ class JumpStartEstimatorInitKwargs(JumpStartKwargs):
         "model_version",
         "hub_arn",
         "model_type",
+        "hub_content_type",
     }
 
     def __init__(

src/sagemaker/jumpstart/utils.py

@@ -384,12 +384,12 @@ def add_jumpstart_model_id_version_tags(
 
 def add_hub_content_arn_tags(
     tags: Optional[List[TagsDict]],
-    hub_arn: str,
+    hub_content_arn: str,
 ) -> Optional[List[TagsDict]]:
     """Adds custom Hub arn tag to JumpStart related resources."""
 
     tags = add_single_jumpstart_tag(
-        hub_arn,
+        hub_content_arn,
         enums.JumpStartTag.HUB_CONTENT_ARN,
         tags,
         is_uri=False,
@@ -1012,24 +1012,37 @@ def get_jumpstart_configs(
     )
 
 
-def get_jumpstart_user_agent_extra_suffix(model_id: str, model_version: str) -> str:
+def get_jumpstart_user_agent_extra_suffix(
+    model_id: Optional[str], model_version: Optional[str], is_hub_content: Optional[bool]
+) -> str:
     """Returns the model-specific user agent string to be added to requests."""
     sagemaker_python_sdk_headers = get_user_agent_extra_suffix()
     jumpstart_specific_suffix = f"md/js_model_id#{model_id} md/js_model_ver#{model_version}"
-    return (
-        sagemaker_python_sdk_headers
-        if os.getenv(constants.ENV_VARIABLE_DISABLE_JUMPSTART_TELEMETRY, None)
-        else f"{sagemaker_python_sdk_headers} {jumpstart_specific_suffix}"
-    )
+    hub_specific_suffix = f"md/js_is_hub_content#{is_hub_content}"
+
+    if os.getenv(constants.ENV_VARIABLE_DISABLE_JUMPSTART_TELEMETRY, None):
+        headers = sagemaker_python_sdk_headers
+    elif model_id is None and model_version is None:
+        headers = f"{sagemaker_python_sdk_headers} {hub_specific_suffix}"
+    else:
+        headers = (
+            f"{sagemaker_python_sdk_headers} {jumpstart_specific_suffix} {hub_specific_suffix}"
+        )
+
+    return headers
 
 
 def get_default_jumpstart_session_with_user_agent_suffix(
-    model_id: str, model_version: str
+    model_id: Optional[str] = None,
+    model_version: Optional[str] = None,
+    is_hub_content: Optional[bool] = False,
 ) -> Session:
     """Returns default JumpStart SageMaker Session with model-specific user agent suffix."""
     botocore_session = botocore.session.get_session()
     botocore_config = botocore.config.Config(
-        user_agent_extra=get_jumpstart_user_agent_extra_suffix(model_id, model_version),
+        user_agent_extra=get_jumpstart_user_agent_extra_suffix(
+            model_id, model_version, is_hub_content
+        ),
     )
     botocore_session.set_default_client_config(botocore_config)
     # shallow copy to not affect default session constant

tests/unit/sagemaker/jumpstart/model/test_model.py

@@ -1360,6 +1360,7 @@ def test_attach(
             model_id="model-id",
             model_version="model-version",
             sagemaker_session=DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+            hub_name=None,
         )
         assert isinstance(val, JumpStartModel)
 

tests/unit/sagemaker/jumpstart/test_utils.py

@@ -1714,21 +1714,21 @@ class TestUserAgent:
     @patch("sagemaker.jumpstart.utils.os.getenv")
     def test_get_jumpstart_user_agent_extra_suffix(self, mock_getenv):
         mock_getenv.return_value = False
-        assert utils.get_jumpstart_user_agent_extra_suffix("some-id", "some-version").endswith(
-            "md/js_model_id#some-id md/js_model_ver#some-version"
-        )
+        assert utils.get_jumpstart_user_agent_extra_suffix(
+            "some-id", "some-version", "False"
+        ).endswith("md/js_model_id#some-id md/js_model_ver#some-version md/js_is_hub_content#False")
         mock_getenv.return_value = None
-        assert utils.get_jumpstart_user_agent_extra_suffix("some-id", "some-version").endswith(
-            "md/js_model_id#some-id md/js_model_ver#some-version"
-        )
+        assert utils.get_jumpstart_user_agent_extra_suffix(
+            "some-id", "some-version", "False"
+        ).endswith("md/js_model_id#some-id md/js_model_ver#some-version md/js_is_hub_content#False")
         mock_getenv.return_value = "True"
-        assert not utils.get_jumpstart_user_agent_extra_suffix("some-id", "some-version").endswith(
-            "md/js_model_id#some-id md/js_model_ver#some-version"
-        )
+        assert not utils.get_jumpstart_user_agent_extra_suffix(
+            "some-id", "some-version", "True"
+        ).endswith("md/js_model_id#some-id md/js_model_ver#some-version md/js_is_hub_content#True")
         mock_getenv.return_value = True
-        assert not utils.get_jumpstart_user_agent_extra_suffix("some-id", "some-version").endswith(
-            "md/js_model_id#some-id md/js_model_ver#some-version"
-        )
+        assert not utils.get_jumpstart_user_agent_extra_suffix(
+            "some-id", "some-version", "True"
+        ).endswith("md/js_model_id#some-id md/js_model_ver#some-version md/js_is_hub_content#True")
 
     @patch("sagemaker.jumpstart.utils.botocore.session")
     @patch("sagemaker.jumpstart.utils.botocore.config.Config")
@@ -1748,7 +1748,7 @@ def test_get_default_jumpstart_session_with_user_agent_suffix(
         utils.get_default_jumpstart_session_with_user_agent_suffix("model_id", "model_version")
         mock_boto3_session.get_session.assert_called_once_with()
         mock_get_jumpstart_user_agent_extra_suffix.assert_called_once_with(
-            "model_id", "model_version"
+            "model_id", "model_version", False
         )
         mock_botocore_config.assert_called_once_with(
             user_agent_extra=mock_get_jumpstart_user_agent_extra_suffix.return_value