fix: Address nits and improvements

Author: Malav Shastri

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,17 @@ def get_model_specs(
                 model_specs.set_hub_content_type(HubContentType.MODEL_REFERENCE)
                 return model_specs
 
+            except Exception as ex:
+                logging.info("Recieved expection while calling APIs for ContentType Model: " + 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/model.py

@@ -34,7 +34,7 @@
     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
@@ -550,7 +550,19 @@ 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/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,