add hub_arn support for accept_types, content_types, serializers, deserializers, and predictor

Author: bencrabtree

src/sagemaker/accept_types.py

@@ -72,6 +72,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,
@@ -85,6 +86,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,6 +113,7 @@ def retrieve_default(
     return artifacts._retrieve_default_accept_type(
         model_id,
         model_version,
+        hub_arn,
         region,
         tolerate_vulnerable_model,
         tolerate_deprecated_model,

src/sagemaker/content_types.py

@@ -72,6 +72,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,
@@ -85,6 +86,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,6 +113,7 @@ def retrieve_default(
     return artifacts._retrieve_default_content_type(
         model_id,
         model_version,
+        hub_arn,
         region,
         tolerate_vulnerable_model,
         tolerate_deprecated_model,

src/sagemaker/deserializers.py

@@ -92,6 +92,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,
@@ -105,6 +106,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,6 +134,7 @@ def retrieve_default(
     return artifacts._retrieve_default_deserializer(
         model_id,
         model_version,
+        hub_arn,
         region,
         tolerate_vulnerable_model,
         tolerate_deprecated_model,

src/sagemaker/jumpstart/artifacts/kwargs.py

@@ -31,6 +31,7 @@
 def _retrieve_model_init_kwargs(
     model_id: str,
     model_version: str,
+    hub_arn: Optional[str] = None,
     region: Optional[str] = None,
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
@@ -43,6 +44,8 @@ def _retrieve_model_init_kwargs(
             retrieve the kwargs.
         model_version (str): Version of the JumpStart model for which to retrieve the
             kwargs.
+        hub_arn (str): The arn of the SageMaker Hub for which to retrieve
+            model details from. (default: None).
         region (Optional[str]): Region for which to retrieve kwargs.
             (Default: None).
         tolerate_vulnerable_model (bool): True if vulnerable versions of model
@@ -66,6 +69,7 @@ def _retrieve_model_init_kwargs(
     model_specs = verify_model_region_and_return_specs(
         model_id=model_id,
         version=model_version,
+        hub_arn=hub_arn,
         scope=JumpStartScriptScope.INFERENCE,
         region=region,
         tolerate_vulnerable_model=tolerate_vulnerable_model,
@@ -85,6 +89,7 @@ def _retrieve_model_deploy_kwargs(
     model_id: str,
     model_version: str,
     instance_type: str,
+    hub_arn: Optional[str] = None,
     region: Optional[str] = None,
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
@@ -99,6 +104,8 @@ def _retrieve_model_deploy_kwargs(
             kwargs.
         instance_type (str): Instance type of the hosting endpoint, to determine if volume size
             is supported.
+        hub_arn (str): The arn of the SageMaker Hub for which to retrieve
+            model details from. (default: None).
         region (Optional[str]): Region for which to retrieve kwargs.
             (Default: None).
         tolerate_vulnerable_model (bool): True if vulnerable versions of model
@@ -123,6 +130,7 @@ def _retrieve_model_deploy_kwargs(
     model_specs = verify_model_region_and_return_specs(
         model_id=model_id,
         version=model_version,
+        hub_arn=hub_arn,
         scope=JumpStartScriptScope.INFERENCE,
         region=region,
         tolerate_vulnerable_model=tolerate_vulnerable_model,

src/sagemaker/jumpstart/artifacts/predictors.py

@@ -72,6 +72,7 @@ def _retrieve_deserializer_from_accept_type(
 def _retrieve_default_deserializer(
     model_id: str,
     model_version: str,
+    hub_arn: Optional[str],
     region: Optional[str],
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
@@ -84,6 +85,8 @@ def _retrieve_default_deserializer(
             retrieve the default deserializer.
         model_version (str): Version of the JumpStart model for which to retrieve the
             default deserializer.
+        hub_arn (str): The arn of the SageMaker Hub for which to retrieve
+            model details from. (Default: None).
         region (Optional[str]): Region for which to retrieve default deserializer.
         tolerate_vulnerable_model (bool): True if vulnerable versions of model
             specifications should be tolerated (exception not raised). If False, raises an
@@ -104,6 +107,7 @@ def _retrieve_default_deserializer(
     default_accept_type = _retrieve_default_accept_type(
         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,
@@ -116,6 +120,7 @@ def _retrieve_default_deserializer(
 def _retrieve_default_serializer(
     model_id: str,
     model_version: str,
+    hub_arn: Optional[str],
     region: Optional[str],
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
@@ -128,6 +133,8 @@ def _retrieve_default_serializer(
             retrieve the default serializer.
         model_version (str): Version of the JumpStart model for which to retrieve the
             default serializer.
+        hub_arn (str): The arn of the SageMaker Hub for which to retrieve
+            model details from. (Default: None).
         region (Optional[str]): Region for which to retrieve default serializer.
         tolerate_vulnerable_model (bool): True if vulnerable versions of model
             specifications should be tolerated (exception not raised). If False, raises an
@@ -147,6 +154,7 @@ def _retrieve_default_serializer(
     default_content_type = _retrieve_default_content_type(
         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,
@@ -273,6 +281,7 @@ def _retrieve_serializer_options(
 def _retrieve_default_content_type(
     model_id: str,
     model_version: str,
+    hub_arn: Optional[str],
     region: Optional[str],
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
@@ -285,6 +294,8 @@ def _retrieve_default_content_type(
             retrieve the default content type.
         model_version (str): Version of the JumpStart model for which to retrieve the
             default content type.
+        hub_arn (str): The arn of the SageMaker Hub for which to retrieve
+            model details from. (default: None).
         region (Optional[str]): Region for which to retrieve default content type.
         tolerate_vulnerable_model (bool): True if vulnerable versions of model
             specifications should be tolerated (exception not raised). If False, raises an
@@ -307,6 +318,7 @@ def _retrieve_default_content_type(
     model_specs = verify_model_region_and_return_specs(
         model_id=model_id,
         version=model_version,
+        hub_arn=hub_arn,
         scope=JumpStartScriptScope.INFERENCE,
         region=region,
         tolerate_vulnerable_model=tolerate_vulnerable_model,
@@ -321,6 +333,7 @@ def _retrieve_default_content_type(
 def _retrieve_default_accept_type(
     model_id: str,
     model_version: str,
+    hub_arn: Optional[str],
     region: Optional[str],
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
@@ -333,6 +346,8 @@ def _retrieve_default_accept_type(
             retrieve the default accept type.
         model_version (str): Version of the JumpStart model for which to retrieve the
             default accept type.
+        hub_arn (str): The arn of the SageMaker Hub for which to retrieve
+            model details from. (Default: None).
         region (Optional[str]): Region for which to retrieve default accept type.
         tolerate_vulnerable_model (bool): True if vulnerable versions of model
             specifications should be tolerated (exception not raised). If False, raises an
@@ -355,6 +370,7 @@ def _retrieve_default_accept_type(
     model_specs = verify_model_region_and_return_specs(
         model_id=model_id,
         version=model_version,
+        hub_arn=hub_arn,
         scope=JumpStartScriptScope.INFERENCE,
         region=region,
         tolerate_vulnerable_model=tolerate_vulnerable_model,

src/sagemaker/jumpstart/artifacts/script_uris.py

@@ -107,6 +107,7 @@ def _retrieve_script_uri(
 def _model_supports_inference_script_uri(
     model_id: str,
     model_version: str,
+    hub_arn: Optional[str],
     region: Optional[str],
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
@@ -119,6 +120,8 @@ def _model_supports_inference_script_uri(
             retrieve the support status for script uri with inference.
         model_version (str): Version of the JumpStart model for which to retrieve the
             support status for script uri with inference.
+        hub_arn (str): The arn of the SageMaker Hub for which to retrieve
+            model details from. (Default: None).
         region (Optional[str]): Region for which to retrieve the
             support status for script uri with inference.
         tolerate_vulnerable_model (bool): True if vulnerable versions of model
@@ -142,6 +145,7 @@ def _model_supports_inference_script_uri(
     model_specs = verify_model_region_and_return_specs(
         model_id=model_id,
         version=model_version,
+        hub_arn=hub_arn,
         scope=JumpStartScriptScope.INFERENCE,
         region=region,
         tolerate_vulnerable_model=tolerate_vulnerable_model,

src/sagemaker/jumpstart/curated_hub/utils.py

@@ -90,15 +90,15 @@ def construct_hub_model_arn_from_inputs(hub_arn: str, model_name: str, version:
 
 
 # TODO: Update to recognize JumpStartHub hub_name
-def generate_hub_arn_for_estimator_init_kwargs(
+def generate_hub_arn_for_init_kwargs(
     hub_name: str, region: Optional[str] = None, session: Optional[Session] = None
 ):
-    """Generates the Hub Arn for JumpStartEstimator from a HubName or Arn.
+    """Generates the Hub Arn for JumpStart class args from a HubName or Arn.
 
     Args:
-        hub_name (str): HubName or HubArn from JumpStartEstimator args
-        region (str): Region from JumpStartEstimator args
-        session (Session): Custom SageMaker Session from JumpStartEstimator args
+        hub_name (str): HubName or HubArn from JumpStart class args
+        region (str): Region from JumpStart class args
+        session (Session): Custom SageMaker Session from JumpStart class args
     """
 
     hub_arn = None

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_estimator_init_kwargs
+from sagemaker.jumpstart.curated_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
 
@@ -523,7 +523,7 @@ def _is_valid_model_id_hook():
 
         hub_arn = None
         if hub_name:
-            hub_arn = generate_hub_arn_for_estimator_init_kwargs(
+            hub_arn = generate_hub_arn_for_init_kwargs(
                 hub_name=hub_name, region=region, session=sagemaker_session
             )
 
@@ -1081,6 +1081,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,