feat: implement curated hub parser and bug bash fixes (aws#1457)

* implement HubContentDocument parser

* modify the parser to remove aliases for hubcontent documents

* bug fix

* update boto3

* Bug Fix in the parser

* Improve Hub Class and related functionalities

* Bug Fix and parser updates

* add missing hub_arn support

* Add model reference deployment support and other minor bug fixes

* fix: retrieve correct image_uri (parser update)

* fix: retrieve correct model URI and model data path from HubContentDocument (parser update)

* Add model reference deployment support

* Model accessor and cache retrival bug fixes

* fix: curated hub model training workflow

* fix: pass sagemaker sessions object to retrieve model specs from describe_hub_content call

* fix: fix payload retrieval for curated hub models

* modify constants, enums

* fix: update parser

* Address nits in the parser

* Add unit tests for parser

* implement pagination for list_models utility

* feat: support wildcard chars for model versions

* Address nits and comments

* Add Hub Content Arn Tag to training and hosting

* Add Hub Content Arn Tag to training and hosting

* fix: HubContentDocument schema version

* fix broken unit tests

* fix prepare_container_def unit tests to include ModelReferenceArn

* fix unit tests for test_session.py

* revert boto version changes

* Fix unit tests

* support wildcard model versions for training workflow

* Add test cases for get_model_versions

* Add/fix unit tests

---------

Co-authored-by: Malav Shastri <[email protected]>

Author: malav-shastri

src/sagemaker/accept_types.py

@@ -123,4 +123,5 @@ def retrieve_default(
         region=region,
         tolerate_vulnerable_model=tolerate_vulnerable_model,
         tolerate_deprecated_model=tolerate_deprecated_model,
+        sagemaker_session=sagemaker_session
     )

src/sagemaker/chainer/model.py

@@ -274,6 +274,7 @@ def prepare_container_def(
         accelerator_type=None,
         serverless_inference_config=None,
         accept_eula=None,
+        model_reference_arn=None
     ):
         """Return a container definition with framework configuration set in model environment.
 
@@ -325,6 +326,7 @@ def prepare_container_def(
             self.model_data,
             deploy_env,
             accept_eula=accept_eula,
+            model_reference_arn=model_reference_arn
         )
 
     def serving_image_uri(

src/sagemaker/djl_inference/model.py

@@ -732,6 +732,7 @@ def prepare_container_def(
         accelerator_type=None,
         serverless_inference_config=None,
         accept_eula=None,
+        model_reference_arn=None
     ):  # pylint: disable=unused-argument
         """A container definition with framework configuration set in model environment variables.
 

src/sagemaker/huggingface/model.py

@@ -470,6 +470,7 @@ def prepare_container_def(
         serverless_inference_config=None,
         inference_tool=None,
         accept_eula=None,
+        model_reference_arn=None
     ):
         """A container definition with framework configuration set in model environment variables.
 
@@ -524,6 +525,7 @@ def prepare_container_def(
             self.repacked_model_data or self.model_data,
             deploy_env,
             accept_eula=accept_eula,
+            model_reference_arn=model_reference_arn
         )
 
     def serving_image_uri(

src/sagemaker/jumpstart/accessors.py

@@ -22,6 +22,8 @@
 from sagemaker.jumpstart import cache
 from sagemaker.jumpstart.hub.utils import construct_hub_model_arn_from_inputs, construct_hub_model_reference_arn_from_inputs
 from sagemaker.jumpstart.constants import JUMPSTART_DEFAULT_REGION_NAME
+from sagemaker.session import Session
+from sagemaker.jumpstart import constants
 
 
 class SageMakerSettings(object):
@@ -257,6 +259,7 @@ def get_model_specs(
         hub_arn: Optional[str] = None,
         s3_client: Optional[boto3.client] = None,
         model_type=JumpStartModelType.OPEN_WEIGHTS,
+        sagemaker_session: Session = constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
     ) -> JumpStartModelSpecs:
         """Returns model specs from JumpStart models cache.
 
@@ -272,6 +275,9 @@ def get_model_specs(
         if s3_client is not None:
             additional_kwargs.update({"s3_client": s3_client})
 
+        if hub_arn:
+            additional_kwargs.update({"sagemaker_session": sagemaker_session})
+
         cache_kwargs = JumpStartModelsAccessor._validate_and_mutate_region_cache_kwargs(
             {**JumpStartModelsAccessor._cache_kwargs, **additional_kwargs}
         )
@@ -282,12 +288,16 @@ def get_model_specs(
                 hub_model_arn = construct_hub_model_arn_from_inputs(
                     hub_arn=hub_arn, model_name=model_id, version=version
                 )
-                return JumpStartModelsAccessor._cache.get_hub_model(hub_model_arn=hub_model_arn)
+                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:
                 hub_model_arn = construct_hub_model_reference_arn_from_inputs(
                     hub_arn=hub_arn, model_name=model_id, version=version
                 )
-                return JumpStartModelsAccessor._cache.get_hub_model_reference(hub_model_arn=hub_model_arn)
+                model_specs = JumpStartModelsAccessor._cache.get_hub_model_reference(hub_model_reference_arn=hub_model_arn)
+                model_specs.set_hub_content_type(HubContentType.MODEL_REFERENCE)
+                return model_specs
 
         return JumpStartModelsAccessor._cache.get_specs(  # type: ignore
             model_id=model_id, version_str=version, model_type=model_type

src/sagemaker/jumpstart/artifacts/image_uris.py

@@ -130,7 +130,11 @@ def _retrieve_image_uri(
             )
             if image_uri is not None:
                 return image_uri
-        ecr_specs = model_specs.hosting_ecr_specs
+        if hub_arn:
+            ecr_uri = model_specs.hosting_ecr_uri
+            return ecr_uri
+        else:
+            ecr_specs = model_specs.hosting_ecr_specs
         if ecr_specs is None:
             raise ValueError(
                 f"No inference ECR configuration found for JumpStart model ID '{model_id}' "
@@ -145,7 +149,11 @@ def _retrieve_image_uri(
             )
             if image_uri is not None:
                 return image_uri
-        ecr_specs = model_specs.training_ecr_specs
+        if hub_arn:
+            ecr_uri = model_specs.training_ecr_uri
+            return ecr_uri
+        else:
+            ecr_specs = model_specs.training_ecr_specs
         if ecr_specs is None:
             raise ValueError(
                 f"No training ECR configuration found for JumpStart model ID '{model_id}' "
@@ -198,6 +206,7 @@ def _retrieve_image_uri(
         version=version_override or ecr_specs.framework_version,
         py_version=ecr_specs.py_version,
         instance_type=instance_type,
+        hub_arn=hub_arn,
         accelerator_type=accelerator_type,
         image_scope=image_scope,
         container_version=container_version,

src/sagemaker/jumpstart/artifacts/model_uris.py

@@ -153,11 +153,15 @@ def _retrieve_model_uri(
 
         is_prepacked = not model_specs.use_inference_script_uri()
 
-        model_artifact_key = (
-            _retrieve_hosting_prepacked_artifact_key(model_specs, instance_type)
-            if is_prepacked
-            else _retrieve_hosting_artifact_key(model_specs, instance_type)
-        )
+        if hub_arn:
+            model_artifact_uri = model_specs.hosting_artifact_uri
+            return model_artifact_uri
+        else:
+            model_artifact_key = (
+                _retrieve_hosting_prepacked_artifact_key(model_specs, instance_type)
+                if is_prepacked
+                else _retrieve_hosting_artifact_key(model_specs, instance_type)
+                )
 
     elif model_scope == JumpStartScriptScope.TRAINING:
 

src/sagemaker/jumpstart/cache.py

@@ -22,6 +22,7 @@
 from packaging.version import Version
 from packaging.specifiers import SpecifierSet, InvalidSpecifier
 from sagemaker.jumpstart.constants import (
+    DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
     ENV_VARIABLE_JUMPSTART_MANIFEST_LOCAL_ROOT_DIR_OVERRIDE,
     ENV_VARIABLE_JUMPSTART_SPECS_LOCAL_ROOT_DIR_OVERRIDE,
     JUMPSTART_DEFAULT_MANIFEST_FILE_S3_KEY,
@@ -48,17 +49,20 @@
     JumpStartModelSpecs,
     JumpStartS3FileType,
     JumpStartVersionedModelId,
-    HubType,
     HubContentType
 )
 from sagemaker.jumpstart.hub import utils as hub_utils
 from sagemaker.jumpstart.hub.interfaces import (
     DescribeHubResponse,
     DescribeHubContentResponse,
 )
+from sagemaker.jumpstart.hub.parsers import (
+    make_model_specs_from_describe_hub_content_response,
+)
 from sagemaker.jumpstart.enums import JumpStartModelType
 from sagemaker.jumpstart import utils
 from sagemaker.utilities.cache import LRUCache
+from sagemaker.session import Session
 
 
 class JumpStartModelsCache:
@@ -84,6 +88,7 @@ def __init__(
         s3_bucket_name: Optional[str] = None,
         s3_client_config: Optional[botocore.config.Config] = None,
         s3_client: Optional[boto3.client] = None,
+        sagemaker_session: Optional[Session] = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
     ) -> None:
         """Initialize a ``JumpStartModelsCache`` instance.
 
@@ -105,6 +110,8 @@ def __init__(
             s3_client_config (Optional[botocore.config.Config]): s3 client config to use for cache.
                 Default: None (no config).
             s3_client (Optional[boto3.client]): s3 client to use. Default: None.
+            sagemaker_session: sagemaker session object to use. 
+                Default: session object from default region us-west-2.
         """
 
         self._region = region or utils.get_region_fallback(
@@ -146,6 +153,7 @@ def __init__(
             if s3_client_config
             else boto3.client("s3", region_name=self._region)
         )
+        self._sagemaker_session = sagemaker_session
 
     def set_region(self, region: str) -> None:
         """Set region for cache. Clears cache after new region is set."""
@@ -453,22 +461,34 @@ def _retrieval_function(
             hub_notebook_description = DescribeHubContentResponse(response)
             return JumpStartCachedContentValue(formatted_content=hub_notebook_description)
 
-        if data_type in [HubContentType.MODEL, HubContentType.MODEL_REFERENCE]:
-            hub_name, _, model_name, model_version = hub_utils.get_info_from_hub_resource_arn(
+        if data_type in {
+            HubContentType.MODEL, 
+            HubContentType.MODEL_REFERENCE,
+        }:
+            
+            hub_arn_extracted_info = hub_utils.get_info_from_hub_resource_arn(
                 id_info
             )
+        
+            model_version: str = hub_utils.get_hub_model_version(
+                hub_model_name=hub_arn_extracted_info.hub_content_name,
+                hub_model_type=data_type.value,
+                hub_name=hub_arn_extracted_info.hub_name,
+                sagemaker_session=self._sagemaker_session,
+                hub_model_version=hub_arn_extracted_info.hub_content_version
+            )
+            
             hub_model_description: Dict[str, Any] = self._sagemaker_session.describe_hub_content(
-                hub_name=hub_name,
-                hub_content_name=model_name,
+                hub_name=hub_arn_extracted_info.hub_name,
+                hub_content_name=hub_arn_extracted_info.hub_content_name,
                 hub_content_version=model_version,
-                hub_content_type=data_type,
+                hub_content_type=data_type.value,
             )
 
             model_specs = make_model_specs_from_describe_hub_content_response(
                 DescribeHubContentResponse(hub_model_description),
             )
 
-            utils.emit_logs_based_on_model_specs(model_specs, self.get_region(), self._s3_client)
             return JumpStartCachedContentValue(formatted_content=model_specs)
 
         raise ValueError(self._file_type_error_msg(data_type))

src/sagemaker/jumpstart/constants.py

@@ -184,6 +184,8 @@
 
 JUMPSTART_DEFAULT_REGION_NAME = boto3.session.Session().region_name or "us-west-2"
 
+JUMPSTART_MODEL_HUB_NAME = "SageMakerPublicHub"
+
 JUMPSTART_DEFAULT_MANIFEST_FILE_S3_KEY = "models_manifest.json"
 JUMPSTART_DEFAULT_PROPRIETARY_MANIFEST_KEY = "proprietary-sdk-manifest.json"
 

src/sagemaker/jumpstart/enums.py

@@ -15,6 +15,7 @@
 from __future__ import absolute_import
 
 from enum import Enum
+from typing import List
 
 
 class ModelFramework(str, Enum):
@@ -93,6 +94,7 @@ class JumpStartTag(str, Enum):
     MODEL_VERSION = "sagemaker-sdk:jumpstart-model-version"
     MODEL_TYPE = "sagemaker-sdk:jumpstart-model-type"
 
+    HUB_CONTENT_ARN = "sagemaker-sdk:hub-content-arn"
 
 class SerializerType(str, Enum):
     """Enum class for serializers associated with JumpStart models."""
@@ -124,6 +126,27 @@ def from_suffixed_type(mime_type_with_suffix: str) -> "MIMEType":
         """Removes suffix from type and instantiates enum."""
         base_type, _, _ = mime_type_with_suffix.partition(";")
         return MIMEType(base_type)
+    
+class NamingConventionType(str, Enum):
+    """Enum class for naming conventions."""
+
+    SNAKE_CASE = "snake_case"
+    UPPER_CAMEL_CASE = "upper_camel_case"
+    DEFAULT = UPPER_CAMEL_CASE
+
+
+class ModelSpecKwargType(str, Enum):
+    """Enum class for types of kwargs for model hub content document and model specs."""
+
+    FIT = "fit_kwargs"
+    MODEL = "model_kwargs"
+    ESTIMATOR = "estimator_kwargs"
+    DEPLOY = "deploy_kwargs"
+
+    @classmethod
+    def arg_keys(cls) -> List[str]:
+        """Returns a list of kwargs keys that each type can have"""
+        return [member.value for member in cls]
 
 
 class JumpStartConfigRankingName(str, Enum):

src/sagemaker/jumpstart/estimator.py

@@ -28,7 +28,7 @@
 from sagemaker.instance_group import InstanceGroup
 from sagemaker.jumpstart.accessors import JumpStartModelsAccessor
 from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
-from sagemaker.jumpstart.curated_hub.utils import generate_hub_arn_for_init_kwargs
+from sagemaker.jumpstart.hub.utils import generate_hub_arn_for_init_kwargs
 from sagemaker.jumpstart.enums import JumpStartScriptScope
 from sagemaker.jumpstart.exceptions import INVALID_MODEL_ID_ERROR_MSG
 
@@ -511,28 +511,29 @@ def __init__(
             ValueError: If the model ID is not recognized by JumpStart.
         """
 
+        hub_arn = None
+        if hub_name:
+            hub_arn = generate_hub_arn_for_init_kwargs(
+                hub_name=hub_name, region=region, session=sagemaker_session
+            )
+        
         def _validate_model_id_and_get_type_hook():
             return validate_model_id_and_get_type(
                 model_id=model_id,
                 model_version=model_version,
                 region=region or getattr(sagemaker_session, "boto_region_name", None),
                 script=JumpStartScriptScope.TRAINING,
                 sagemaker_session=sagemaker_session,
+                hub_arn=hub_arn
             )
-
+        
         self.model_type = _validate_model_id_and_get_type_hook()
         if not self.model_type:
             JumpStartModelsAccessor.reset_cache()
             self.model_type = _validate_model_id_and_get_type_hook()
-            if not self.model_type:
+            if not self.model_type and not hub_arn:
                 raise ValueError(INVALID_MODEL_ID_ERROR_MSG.format(model_id=model_id))
 
-        hub_arn = None
-        if hub_name:
-            hub_arn = generate_hub_arn_for_init_kwargs(
-                hub_name=hub_name, region=region, session=sagemaker_session
-            )
-
         estimator_init_kwargs = get_init_kwargs(
             model_id=model_id,
             model_version=model_version,
@@ -691,6 +692,7 @@ def attach(
         training_job_name: str,
         model_id: Optional[str] = None,
         model_version: Optional[str] = None,
+        hub_arn: Optional[str] = None,
         sagemaker_session: session.Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
         model_channel_name: str = "model",
     ) -> "JumpStartEstimator":
@@ -756,6 +758,7 @@ def attach(
         model_specs = verify_model_region_and_return_specs(
             model_id=model_id,
             version=model_version,
+            hub_arn=hub_arn,
             region=sagemaker_session.boto_region_name,
             scope=JumpStartScriptScope.TRAINING,
             tolerate_deprecated_model=True,  # model is already trained, so tolerate if deprecated
@@ -1110,6 +1113,7 @@ def deploy(
                 predictor=predictor,
                 model_id=self.model_id,
                 model_version=self.model_version,
+                hub_arn=self.hub_arn,
                 region=self.region,
                 tolerate_deprecated_model=self.tolerate_deprecated_model,
                 tolerate_vulnerable_model=self.tolerate_vulnerable_model,