JumpStart CuratedHub Launch (aws#4748)

* Implement CuratedHub APIs (aws#1449)

* Implement CuratedHub Admin APIs

* making some parameters optional in create_hub_content_reference as per the API design

* add describe_hub and list_hubs APIs

* implement delete_hub API

* Implement list_hub_contents API

* create CuratedHub class and supported utils

* implement list_models and address comments

* Add unit tests

* add describe_model function

* cache retrieval for describeHubContent changes

* fix curated hub class unit tests

* add utils needed for curatedHub

* Cache retrieval

* implement get_hub_model_reference()

* cleanup HUB type datatype

* cleanup constants

* rename list_public_models to list_jumpstart_service_hub_models

* implement describe_model_reference

* Rename CuratedHub to Hub

* address nit

* address nits and fix failing tests

---------

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

* feat: implement list_jumpstart_service_hub_models function to fetch JumpStart public hub models (aws#1456)

* Implement CuratedHub Admin APIs

* making some parameters optional in create_hub_content_reference as per the API design

* add describe_hub and list_hubs APIs

* implement delete_hub API

* Implement list_hub_contents API

* create CuratedHub class and supported utils

* implement list_models and address comments

* Add unit tests

* add describe_model function

* cache retrieval for describeHubContent changes

* fix curated hub class unit tests

* add utils needed for curatedHub

* Cache retrieval

* implement get_hub_model_reference()

* cleanup HUB type datatype

* cleanup constants

* rename list_public_models to list_jumpstart_service_hub_models

* implement describe_model_reference

* Rename CuratedHub to Hub

* address nit

* address nits and fix failing tests

* implement list_jumpstart_service_hub_models function

---------

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

* Feat/Curated Hub hub_arn and hub_content_type support (aws#1453)

* get_model_spec() changes to support hub_arn and hub_content_type

* implement get_hub_model_reference()

* support hub_arn and hub_content_type for specs retrieval

* add support for hub_arn and hub_content_type for serializers, deserializers, estimators, models, predictors and various spec retrieval functionalities

* address nits and test failures

* remove hub_content_type support

---------

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

* 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]>

* address unit tests failures in codebuild

* change list_jumpstart_service_hub_models to list_sagemaker_public_hub_models()

* fix: Changing list input output shapes

* fix: gated model training bug

* run black -l 100

* flake 8

* address formatting issues

* black -l

* DocStyle issues

* address flake8, pylint

* blake -l

* pass model type down

* disabling pylint for release

* disable pylint

---------

Co-authored-by: Malav Shastri <[email protected]>
Co-authored-by: chrstfu <[email protected]>
Co-authored-by: Erick Benitez-Ramos <[email protected]>

Author: 4 people

src/sagemaker/accept_types.py

@@ -24,6 +24,7 @@ def retrieve_options(
     region: Optional[str] = None,
     model_id: Optional[str] = None,
     model_version: Optional[str] = None,
+    hub_arn: Optional[str] = None,
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
@@ -37,6 +38,8 @@ def retrieve_options(
             retrieve the supported accept types. (Default: None).
         model_version (str): The version of the model for which to retrieve the
             supported accept types. (Default: None).
+        hub_arn (str): The arn of the SageMaker Hub for which to retrieve
+            model details from. (Default: None).
         tolerate_vulnerable_model (bool): True if vulnerable versions of model
             specifications should be tolerated (exception not raised). If False, raises an
             exception if the script used by this version of the model has dependencies with known
@@ -60,11 +63,12 @@ def retrieve_options(
         )
 
     return artifacts._retrieve_supported_accept_types(
-        model_id,
-        model_version,
-        region,
-        tolerate_vulnerable_model,
-        tolerate_deprecated_model,
+        model_id=model_id,
+        model_version=model_version,
+        hub_arn=hub_arn,
+        region=region,
+        tolerate_vulnerable_model=tolerate_vulnerable_model,
+        tolerate_deprecated_model=tolerate_deprecated_model,
         sagemaker_session=sagemaker_session,
     )
 
@@ -73,6 +77,7 @@ def retrieve_default(
     region: Optional[str] = None,
     model_id: Optional[str] = None,
     model_version: Optional[str] = None,
+    hub_arn: Optional[str] = None,
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
@@ -87,6 +92,8 @@ def retrieve_default(
             retrieve the default accept type. (Default: None).
         model_version (str): The version of the model for which to retrieve the
             default accept type. (Default: None).
+        hub_arn (str): The arn of the SageMaker Hub for which to retrieve
+            model details from. (Default: None).
         tolerate_vulnerable_model (bool): True if vulnerable versions of model
             specifications should be tolerated (exception not raised). If False, raises an
             exception if the script used by this version of the model has dependencies with known
@@ -110,11 +117,12 @@ def retrieve_default(
         )
 
     return artifacts._retrieve_default_accept_type(
-        model_id,
-        model_version,
-        region,
-        tolerate_vulnerable_model,
-        tolerate_deprecated_model,
+        model_id=model_id,
+        model_version=model_version,
+        hub_arn=hub_arn,
+        region=region,
+        tolerate_vulnerable_model=tolerate_vulnerable_model,
+        tolerate_deprecated_model=tolerate_deprecated_model,
         sagemaker_session=sagemaker_session,
         model_type=model_type,
     )

src/sagemaker/chainer/model.py

@@ -282,6 +282,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.
 
@@ -333,6 +334,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/content_types.py

@@ -24,6 +24,7 @@ def retrieve_options(
     region: Optional[str] = None,
     model_id: Optional[str] = None,
     model_version: Optional[str] = None,
+    hub_arn: Optional[str] = None,
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
@@ -37,6 +38,8 @@ def retrieve_options(
             retrieve the supported content types. (Default: None).
         model_version (str): The version of the model for which to retrieve the
             supported content types. (Default: None).
+        hub_arn (str): The arn of the SageMaker Hub for which to retrieve
+            model details from. (Default: None).
         tolerate_vulnerable_model (bool): True if vulnerable versions of model
             specifications should be tolerated (exception not raised). If False, raises an
             exception if the script used by this version of the model has dependencies with known
@@ -60,11 +63,12 @@ def retrieve_options(
         )
 
     return artifacts._retrieve_supported_content_types(
-        model_id,
-        model_version,
-        region,
-        tolerate_vulnerable_model,
-        tolerate_deprecated_model,
+        model_id=model_id,
+        model_version=model_version,
+        hub_arn=hub_arn,
+        region=region,
+        tolerate_vulnerable_model=tolerate_vulnerable_model,
+        tolerate_deprecated_model=tolerate_deprecated_model,
         sagemaker_session=sagemaker_session,
     )
 
@@ -73,6 +77,7 @@ def retrieve_default(
     region: Optional[str] = None,
     model_id: Optional[str] = None,
     model_version: Optional[str] = None,
+    hub_arn: Optional[str] = None,
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
@@ -87,6 +92,8 @@ def retrieve_default(
             retrieve the default content type. (Default: None).
         model_version (str): The version of the model for which to retrieve the
             default content type. (Default: None).
+        hub_arn (str): The arn of the SageMaker Hub for which to retrieve
+            model details from. (default: None).
         tolerate_vulnerable_model (bool): True if vulnerable versions of model
             specifications should be tolerated (exception not raised). If False, raises an
             exception if the script used by this version of the model has dependencies with known
@@ -110,11 +117,12 @@ def retrieve_default(
         )
 
     return artifacts._retrieve_default_content_type(
-        model_id,
-        model_version,
-        region,
-        tolerate_vulnerable_model,
-        tolerate_deprecated_model,
+        model_id=model_id,
+        model_version=model_version,
+        hub_arn=hub_arn,
+        region=region,
+        tolerate_vulnerable_model=tolerate_vulnerable_model,
+        tolerate_deprecated_model=tolerate_deprecated_model,
         sagemaker_session=sagemaker_session,
         model_type=model_type,
     )

src/sagemaker/deserializers.py

@@ -43,6 +43,7 @@ def retrieve_options(
     region: Optional[str] = None,
     model_id: Optional[str] = None,
     model_version: Optional[str] = None,
+    hub_arn: Optional[str] = None,
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
@@ -56,6 +57,8 @@ def retrieve_options(
             retrieve the supported deserializers. (Default: None).
         model_version (str): The version of the model for which to retrieve the
             supported deserializers. (Default: None).
+        hub_arn (str): The arn of the SageMaker Hub for which to retrieve
+            model details from. (Default: None).
         tolerate_vulnerable_model (bool): True if vulnerable versions of model
             specifications should be tolerated (exception not raised). If False, raises an
             exception if the script used by this version of the model has dependencies with known
@@ -80,11 +83,12 @@ def retrieve_options(
         )
 
     return artifacts._retrieve_deserializer_options(
-        model_id,
-        model_version,
-        region,
-        tolerate_vulnerable_model,
-        tolerate_deprecated_model,
+        model_id=model_id,
+        model_version=model_version,
+        hub_arn=hub_arn,
+        region=region,
+        tolerate_vulnerable_model=tolerate_vulnerable_model,
+        tolerate_deprecated_model=tolerate_deprecated_model,
         sagemaker_session=sagemaker_session,
     )
 
@@ -93,6 +97,7 @@ def retrieve_default(
     region: Optional[str] = None,
     model_id: Optional[str] = None,
     model_version: Optional[str] = None,
+    hub_arn: Optional[str] = None,
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
@@ -107,6 +112,8 @@ def retrieve_default(
             retrieve the default deserializer. (Default: None).
         model_version (str): The version of the model for which to retrieve the
             default deserializer. (Default: None).
+        hub_arn (str): The arn of the SageMaker Hub for which to retrieve
+            model details from. (Default: None).
         tolerate_vulnerable_model (bool): True if vulnerable versions of model
             specifications should be tolerated (exception not raised). If False, raises an
             exception if the script used by this version of the model has dependencies with known
@@ -131,11 +138,12 @@ def retrieve_default(
         )
 
     return artifacts._retrieve_default_deserializer(
-        model_id,
-        model_version,
-        region,
-        tolerate_vulnerable_model,
-        tolerate_deprecated_model,
+        model_id=model_id,
+        model_version=model_version,
+        hub_arn=hub_arn,
+        region=region,
+        tolerate_vulnerable_model=tolerate_vulnerable_model,
+        tolerate_deprecated_model=tolerate_deprecated_model,
         sagemaker_session=sagemaker_session,
         model_type=model_type,
     )

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/environment_variables.py

@@ -30,6 +30,7 @@ def retrieve_default(
     region: Optional[str] = None,
     model_id: Optional[str] = None,
     model_version: Optional[str] = None,
+    hub_arn: Optional[str] = None,
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
     include_aws_sdk_env_vars: bool = True,
@@ -46,6 +47,8 @@ def retrieve_default(
             retrieve the default environment variables. (Default: None).
         model_version (str): Optional. The version of the model for which to retrieve the
             default environment variables. (Default: None).
+        hub_arn (str): The arn of the SageMaker Hub for which to
+            retrieve model details from. (Default: None).
         tolerate_vulnerable_model (bool): True if vulnerable versions of model
             specifications should be tolerated (exception not raised). If False, raises an
             exception if the script used by this version of the model has dependencies with known
@@ -78,12 +81,13 @@ def retrieve_default(
         )
 
     return artifacts._retrieve_default_environment_variables(
-        model_id,
-        model_version,
-        region,
-        tolerate_vulnerable_model,
-        tolerate_deprecated_model,
-        include_aws_sdk_env_vars,
+        model_id=model_id,
+        model_version=model_version,
+        hub_arn=hub_arn,
+        region=region,
+        tolerate_vulnerable_model=tolerate_vulnerable_model,
+        tolerate_deprecated_model=tolerate_deprecated_model,
+        include_aws_sdk_env_vars=include_aws_sdk_env_vars,
         sagemaker_session=sagemaker_session,
         instance_type=instance_type,
         script=script,

src/sagemaker/huggingface/model.py

@@ -479,6 +479,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.
 
@@ -533,6 +534,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/hyperparameters.py

@@ -31,6 +31,7 @@ def retrieve_default(
     region: Optional[str] = None,
     model_id: Optional[str] = None,
     model_version: Optional[str] = None,
+    hub_arn: Optional[str] = None,
     instance_type: Optional[str] = None,
     include_container_hyperparameters: bool = False,
     tolerate_vulnerable_model: bool = False,
@@ -46,6 +47,8 @@ def retrieve_default(
             retrieve the default hyperparameters. (Default: None).
         model_version (str): The version of the model for which to retrieve the
             default hyperparameters. (Default: None).
+        hub_arn (str): The arn of the SageMaker Hub for which to retrieve
+            model details from. (default: None).
         instance_type (str): An instance type to optionally supply in order to get hyperparameters
             specific for the instance type.
         include_container_hyperparameters (bool): ``True`` if the container hyperparameters
@@ -80,6 +83,7 @@ def retrieve_default(
     return artifacts._retrieve_default_hyperparameters(
         model_id=model_id,
         model_version=model_version,
+        hub_arn=hub_arn,
         instance_type=instance_type,
         region=region,
         include_container_hyperparameters=include_container_hyperparameters,
@@ -92,6 +96,7 @@ def retrieve_default(
 def validate(
     region: Optional[str] = None,
     model_id: Optional[str] = None,
+    hub_arn: Optional[str] = None,
     model_version: Optional[str] = None,
     hyperparameters: Optional[dict] = None,
     validation_mode: HyperparameterValidationMode = HyperparameterValidationMode.VALIDATE_PROVIDED,
@@ -107,6 +112,8 @@ def validate(
             (Default: None).
         model_version (str): The version of the model for which to validate hyperparameters.
             (Default: None).
+        hub_arn (str): The arn of the SageMaker Hub for which to retrieve
+            model details from. (default: None).
         hyperparameters (dict): Hyperparameters to validate.
             (Default: None).
         validation_mode (HyperparameterValidationMode): Method of validation to use with
@@ -148,6 +155,7 @@ def validate(
     return validate_hyperparameters(
         model_id=model_id,
         model_version=model_version,
+        hub_arn=hub_arn,
         hyperparameters=hyperparameters,
         validation_mode=validation_mode,
         region=region,

src/sagemaker/image_uris.py

@@ -64,6 +64,7 @@ def retrieve(
     training_compiler_config=None,
     model_id=None,
     model_version=None,
+    hub_arn=None,
     tolerate_vulnerable_model=False,
     tolerate_deprecated_model=False,
     sdk_version=None,
@@ -104,6 +105,8 @@ def retrieve(
             (default: None).
         model_version (str): The version of the JumpStart model for which to retrieve the
             image URI (default: None).
+        hub_arn (str): The arn of the SageMaker Hub for which to retrieve
+            model details from. (Default: None).
         tolerate_vulnerable_model (bool): ``True`` if vulnerable versions of model specifications
             should be tolerated without an exception raised. If ``False``, raises an exception if
             the script used by this version of the model has dependencies with known security
@@ -149,6 +152,7 @@ def retrieve(
             model_id,
             model_version,
             image_scope,
+            hub_arn,
             framework,
             region,
             version,