add hub name support for jumpstart estimator

Author: bencrabtree

src/sagemaker/instance_types.py

@@ -29,6 +29,7 @@ def retrieve_default(
     region: Optional[str] = None,
     model_id: Optional[str] = None,
     model_version: Optional[str] = None,
+    hub_arn: Optional[str] = None,
     scope: Optional[str] = None,
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
@@ -80,6 +81,7 @@ def retrieve_default(
         model_id,
         model_version,
         scope,
+        hub_arn,
         region,
         tolerate_vulnerable_model,
         tolerate_deprecated_model,
@@ -92,6 +94,7 @@ def retrieve(
     region: Optional[str] = None,
     model_id: Optional[str] = None,
     model_version: Optional[str] = None,
+    hub_arn: Optional[str] = None,
     scope: Optional[str] = None,
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
@@ -142,6 +145,7 @@ def retrieve(
         model_id,
         model_version,
         scope,
+        hub_arn,
         region,
         tolerate_vulnerable_model,
         tolerate_deprecated_model,

src/sagemaker/jumpstart/accessors.py

@@ -18,7 +18,7 @@
 
 from sagemaker.deprecations import deprecated
 from sagemaker.jumpstart.types import JumpStartModelHeader, JumpStartModelSpecs
-from sagemaker.jumpstart import cache
+from sagemaker.jumpstart import cache, utils
 from sagemaker.jumpstart.constants import JUMPSTART_DEFAULT_REGION_NAME
 
 
@@ -239,7 +239,11 @@ def get_model_header(region: str, model_id: str, version: str) -> JumpStartModel
 
     @staticmethod
     def get_model_specs(
-        region: str, model_id: str, version: str, s3_client: Optional[boto3.client] = None
+        region: str,
+        model_id: str,
+        version: str,
+        hub_arn: Optional[str] = None,
+        s3_client: Optional[boto3.client] = None,
     ) -> JumpStartModelSpecs:
         """Returns model specs from JumpStart models cache.
 
@@ -259,6 +263,13 @@ def get_model_specs(
             {**JumpStartModelsAccessor._cache_kwargs, **additional_kwargs}
         )
         JumpStartModelsAccessor._set_cache_and_region(region, cache_kwargs)
+
+        if hub_arn:
+            hub_model_arn = utils.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)
+
         return JumpStartModelsAccessor._cache.get_specs(  # type: ignore
             model_id=model_id, semantic_version_str=version
         )

src/sagemaker/jumpstart/artifacts/instance_types.py

@@ -33,6 +33,7 @@ def _retrieve_default_instance_type(
     model_id: str,
     model_version: str,
     scope: str,
+    hub_arn: Optional[str] = None,
     region: Optional[str] = None,
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
@@ -80,6 +81,7 @@ def _retrieve_default_instance_type(
     model_specs = verify_model_region_and_return_specs(
         model_id=model_id,
         version=model_version,
+        hub_arn=hub_arn,
         scope=scope,
         region=region,
         tolerate_vulnerable_model=tolerate_vulnerable_model,
@@ -119,6 +121,7 @@ def _retrieve_instance_types(
     model_id: str,
     model_version: str,
     scope: str,
+    hub_arn: Optional[str] = None,
     region: Optional[str] = None,
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
@@ -166,6 +169,7 @@ def _retrieve_instance_types(
     model_specs = verify_model_region_and_return_specs(
         model_id=model_id,
         version=model_version,
+        hub_arn=hub_arn,
         scope=scope,
         region=region,
         tolerate_vulnerable_model=tolerate_vulnerable_model,

src/sagemaker/jumpstart/artifacts/kwargs.py

@@ -198,6 +198,7 @@ def _retrieve_estimator_init_kwargs(
 def _retrieve_estimator_fit_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,
@@ -234,6 +235,7 @@ def _retrieve_estimator_fit_kwargs(
     model_specs = verify_model_region_and_return_specs(
         model_id=model_id,
         version=model_version,
+        hub_arn=hub_arn,
         scope=JumpStartScriptScope.TRAINING,
         region=region,
         tolerate_vulnerable_model=tolerate_vulnerable_model,

src/sagemaker/jumpstart/constants.py

@@ -170,8 +170,8 @@
 
 JUMPSTART_DEFAULT_MANIFEST_FILE_S3_KEY = "models_manifest.json"
 
-# works cross-partition
-HUB_MODEL_ARN_REGEX = r"arn:(.*?):sagemaker:(.*?):(.*?):hub_content/(.*?)/Model/(.*?)/(.*?)$"
+# works for cross-partition
+HUB_CONTENT_ARN_REGEX = r"arn:(.*?):sagemaker:(.*?):(.*?):hub_content/(.*?)/(.*?)/(.*?)/(.*?)$"
 HUB_ARN_REGEX = r"arn:(.*?):sagemaker:(.*?):(.*?):hub/(.*?)$"
 
 INFERENCE_ENTRY_POINT_SCRIPT_NAME = "inference.py"

src/sagemaker/jumpstart/curated_hub/curated_hub.py

@@ -25,7 +25,7 @@ def __init__(self, hub_name: str, region: str, session: Optional[Session] = None
         self.hub_name = hub_name
         self.region = region
         self.session = session
-        self._sm_session = session or Session()
+        self._sm_session = session or DEFAULT_JUMPSTART_SAGEMAKER_SESSION
 
     def describe_model(self, model_name: str, model_version: str = "*") -> Dict[str, Any]:
         """Returns descriptive information about the Hub Model"""

src/sagemaker/jumpstart/enums.py

@@ -79,6 +79,8 @@ class JumpStartTag(str, Enum):
     MODEL_ID = "sagemaker-sdk:jumpstart-model-id"
     MODEL_VERSION = "sagemaker-sdk:jumpstart-model-version"
 
+    HUB_ARN = "sagemaker-hub:hub-arn"
+
 
 class SerializerType(str, Enum):
     """Enum class for serializers associated with JumpStart models."""

src/sagemaker/jumpstart/estimator.py

@@ -12,6 +12,7 @@
 # language governing permissions and limitations under the License.
 """This module stores JumpStart implementation of Estimator class."""
 from __future__ import absolute_import
+import re
 
 
 from typing import Dict, List, Optional, Union
@@ -27,14 +28,15 @@
 from sagemaker.inputs import FileSystemInput, TrainingInput
 from sagemaker.instance_group import InstanceGroup
 from sagemaker.jumpstart.accessors import JumpStartModelsAccessor
-from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
+from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION, HUB_ARN_REGEX
 from sagemaker.jumpstart.enums import JumpStartScriptScope
 from sagemaker.jumpstart.exceptions import INVALID_MODEL_ID_ERROR_MSG
 
 from sagemaker.jumpstart.factory.estimator import get_deploy_kwargs, get_fit_kwargs, get_init_kwargs
 from sagemaker.jumpstart.factory.model import get_default_predictor
 from sagemaker.jumpstart.session_utils import get_model_id_version_from_training_job
 from sagemaker.jumpstart.utils import (
+    construct_hub_arn_from_name,
     is_valid_model_id,
     resolve_model_sagemaker_config_field,
 )
@@ -57,6 +59,7 @@ def __init__(
         self,
         model_id: Optional[str] = None,
         model_version: Optional[str] = None,
+        hub_name: Optional[str] = None,
         tolerate_vulnerable_model: Optional[bool] = None,
         tolerate_deprecated_model: Optional[bool] = None,
         region: Optional[str] = None,
@@ -122,6 +125,7 @@ def __init__(
                 https://sagemaker.readthedocs.io/en/stable/doc_utils/pretrainedmodels.html
                 for list of model IDs.
             model_version (Optional[str]): Version for JumpStart model to use (Default: None).
+            hub_name (Optional[str]): Hub name or arn where the model is stored (Default: None).
             tolerate_vulnerable_model (Optional[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
@@ -518,9 +522,19 @@ def _is_valid_model_id_hook():
             if not _is_valid_model_id_hook():
                 raise ValueError(INVALID_MODEL_ID_ERROR_MSG.format(model_id=model_id))
 
+        # TODO: Update to handle SageMakerJumpStart hub_name
+        hub_arn = None
+        if hub_name:
+            match = re.match(HUB_ARN_REGEX, hub_name)
+            if match:
+                hub_arn = hub_name
+            else:
+                hub_arn = construct_hub_arn_from_name(hub_name=hub_name, region=region, sagemaker_session=sagemaker_session)
+
         estimator_init_kwargs = get_init_kwargs(
             model_id=model_id,
             model_version=model_version,
+            hub_arn=hub_arn,
             tolerate_vulnerable_model=tolerate_vulnerable_model,
             tolerate_deprecated_model=tolerate_deprecated_model,
             role=role,
@@ -576,6 +590,7 @@ def _is_valid_model_id_hook():
             enable_remote_debug=enable_remote_debug,
         )
 
+        self.hub_arn = estimator_init_kwargs.hub_arn
         self.model_id = estimator_init_kwargs.model_id
         self.model_version = estimator_init_kwargs.model_version
         self.instance_type = estimator_init_kwargs.instance_type
@@ -652,6 +667,7 @@ def fit(
         estimator_fit_kwargs = get_fit_kwargs(
             model_id=self.model_id,
             model_version=self.model_version,
+            hub_arn=self.hub_arn,
             region=self.region,
             inputs=inputs,
             wait=wait,
@@ -1018,6 +1034,7 @@ def deploy(
         estimator_deploy_kwargs = get_deploy_kwargs(
             model_id=self.model_id,
             model_version=self.model_version,
+            hub_arn=self.hub_arn,
             region=self.region,
             tolerate_vulnerable_model=self.tolerate_vulnerable_model,
             tolerate_deprecated_model=self.tolerate_deprecated_model,

src/sagemaker/jumpstart/factory/estimator.py

@@ -61,6 +61,7 @@
     JumpStartModelInitKwargs,
 )
 from sagemaker.jumpstart.utils import (
+    add_hub_arn_tags,
     add_jumpstart_model_id_version_tags,
     update_dict_if_key_not_present,
     resolve_estimator_sagemaker_config_field,
@@ -77,6 +78,7 @@
 def get_init_kwargs(
     model_id: str,
     model_version: Optional[str] = None,
+    hub_arn: Optional[str] = None,
     tolerate_vulnerable_model: Optional[bool] = None,
     tolerate_deprecated_model: Optional[bool] = None,
     region: Optional[str] = None,
@@ -134,6 +136,7 @@ def get_init_kwargs(
     estimator_init_kwargs: JumpStartEstimatorInitKwargs = JumpStartEstimatorInitKwargs(
         model_id=model_id,
         model_version=model_version,
+        hub_arn=hub_arn,
         role=role,
         region=region,
         instance_count=instance_count,
@@ -209,6 +212,7 @@ def get_init_kwargs(
 def get_fit_kwargs(
     model_id: str,
     model_version: Optional[str] = None,
+    hub_arn: Optional[str] = None,
     region: Optional[str] = None,
     inputs: Optional[Union[str, Dict, TrainingInput, FileSystemInput]] = None,
     wait: Optional[bool] = None,
@@ -224,6 +228,7 @@ def get_fit_kwargs(
     estimator_fit_kwargs: JumpStartEstimatorFitKwargs = JumpStartEstimatorFitKwargs(
         model_id=model_id,
         model_version=model_version,
+        hub_arn=hub_arn,
         region=region,
         inputs=inputs,
         wait=wait,
@@ -246,6 +251,7 @@ def get_fit_kwargs(
 def get_deploy_kwargs(
     model_id: str,
     model_version: Optional[str] = None,
+    hub_arn: Optional[str] = None,
     region: Optional[str] = None,
     initial_instance_count: Optional[int] = None,
     instance_type: Optional[str] = None,
@@ -290,6 +296,7 @@ def get_deploy_kwargs(
     model_deploy_kwargs: JumpStartModelDeployKwargs = model.get_deploy_kwargs(
         model_id=model_id,
         model_version=model_version,
+        hub_arn=hub_arn,
         region=region,
         initial_instance_count=initial_instance_count,
         instance_type=instance_type,
@@ -432,6 +439,7 @@ def _add_instance_type_and_count_to_kwargs(
         region=kwargs.region,
         model_id=kwargs.model_id,
         model_version=kwargs.model_version,
+        hub_arn=kwargs.hub_arn,
         scope=JumpStartScriptScope.TRAINING,
         tolerate_deprecated_model=kwargs.tolerate_deprecated_model,
         tolerate_vulnerable_model=kwargs.tolerate_vulnerable_model,
@@ -465,6 +473,10 @@ def _add_tags_to_kwargs(kwargs: JumpStartEstimatorInitKwargs) -> JumpStartEstima
         kwargs.tags = add_jumpstart_model_id_version_tags(
             kwargs.tags, kwargs.model_id, full_model_version
         )
+
+    if kwargs.hub_arn:
+        kwargs.tags = add_hub_arn_tags(kwargs.tags, kwargs.hub_arn)
+
     return kwargs
 
 
@@ -728,6 +740,7 @@ def _add_fit_extra_kwargs(kwargs: JumpStartEstimatorFitKwargs) -> JumpStartEstim
     fit_kwargs_to_add = _retrieve_estimator_fit_kwargs(
         model_id=kwargs.model_id,
         model_version=kwargs.model_version,
+        hub_arn=kwargs.hub_arn,
         region=kwargs.region,
         tolerate_deprecated_model=kwargs.tolerate_deprecated_model,
         tolerate_vulnerable_model=kwargs.tolerate_vulnerable_model,

src/sagemaker/jumpstart/factory/model.py

@@ -44,6 +44,7 @@
     JumpStartModelRegisterKwargs,
 )
 from sagemaker.jumpstart.utils import (
+    add_hub_arn_tags,
     add_jumpstart_model_id_version_tags,
     update_dict_if_key_not_present,
     resolve_model_sagemaker_config_field,
@@ -447,6 +448,9 @@ def _add_tags_to_kwargs(kwargs: JumpStartModelDeployKwargs) -> Dict[str, Any]:
             kwargs.tags, kwargs.model_id, full_model_version
         )
 
+    if kwargs.hub_arn:
+        kwargs.tags = add_hub_arn_tags(kwargs.tags, kwargs.hub_arn)
+
     return kwargs
 
 
@@ -489,6 +493,7 @@ def _add_resources_to_kwargs(kwargs: JumpStartModelInitKwargs) -> JumpStartModel
 def get_deploy_kwargs(
     model_id: str,
     model_version: Optional[str] = None,
+    hub_arn: Optional[str] = None,
     region: Optional[str] = None,
     initial_instance_count: Optional[int] = None,
     instance_type: Optional[str] = None,
@@ -521,6 +526,7 @@ def get_deploy_kwargs(
     deploy_kwargs: JumpStartModelDeployKwargs = JumpStartModelDeployKwargs(
         model_id=model_id,
         model_version=model_version,
+        hub_arn=hub_arn,
         region=region,
         initial_instance_count=initial_instance_count,
         instance_type=instance_type,