fix: fix _add_tags_to_kwargs to use hub_content_arn instead of hub_arn

Author: Malav Shastri

src/sagemaker/jumpstart/factory/estimator.py

@@ -29,6 +29,7 @@
     _retrieve_model_package_model_artifact_s3_uri,
 )
 from sagemaker.jumpstart.artifacts.resource_names import _retrieve_resource_name_base
+from sagemaker.jumpstart.hub.utils import construct_hub_model_arn_from_inputs, construct_hub_model_reference_arn_from_inputs
 from sagemaker.session import Session
 from sagemaker.async_inference.async_inference_config import AsyncInferenceConfig
 from sagemaker.base_deserializers import BaseDeserializer
@@ -52,6 +53,7 @@
 from sagemaker.jumpstart.enums import JumpStartScriptScope, JumpStartModelType
 from sagemaker.jumpstart.factory import model
 from sagemaker.jumpstart.types import (
+    HubContentType,
     JumpStartEstimatorDeployKwargs,
     JumpStartEstimatorFitKwargs,
     JumpStartEstimatorInitKwargs,
@@ -201,6 +203,11 @@ def get_init_kwargs(
     estimator_init_kwargs = _add_region_to_kwargs(estimator_init_kwargs)
     estimator_init_kwargs = _add_instance_type_and_count_to_kwargs(estimator_init_kwargs)
     estimator_init_kwargs = _add_image_uri_to_kwargs(estimator_init_kwargs)
+    if hub_arn:
+        estimator_init_kwargs = _add_model_reference_arn_to_kwargs(kwargs=estimator_init_kwargs)
+    else:
+        estimator_init_kwargs.model_reference_arn = None
+        estimator_init_kwargs.hub_content_type = None
     estimator_init_kwargs = _add_model_uri_to_kwargs(estimator_init_kwargs)
     estimator_init_kwargs = _add_source_dir_to_kwargs(estimator_init_kwargs)
     estimator_init_kwargs = _add_entry_point_to_kwargs(estimator_init_kwargs)
@@ -511,7 +518,15 @@ def _add_tags_to_kwargs(kwargs: JumpStartEstimatorInitKwargs) -> JumpStartEstima
         )
 
     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
 
@@ -534,6 +549,32 @@ def _add_image_uri_to_kwargs(kwargs: JumpStartEstimatorInitKwargs) -> JumpStartE
 
     return kwargs
 
+def _add_model_reference_arn_to_kwargs(
+    kwargs: JumpStartEstimatorInitKwargs,
+) -> JumpStartEstimatorInitKwargs:
+    """Sets Model Reference ARN if the hub content type is Model Reference, returns full kwargs."""
+    
+    hub_content_type = verify_model_region_and_return_specs(
+        model_id=kwargs.model_id,
+        version=kwargs.model_version,
+        hub_arn=kwargs.hub_arn,
+        scope=JumpStartScriptScope.TRAINING,
+        region=kwargs.region,
+        tolerate_vulnerable_model=kwargs.tolerate_vulnerable_model,
+        tolerate_deprecated_model=kwargs.tolerate_deprecated_model,
+        sagemaker_session=kwargs.sagemaker_session,
+        model_type=kwargs.model_type,
+    ).hub_content_type
+    kwargs.hub_content_type = hub_content_type if kwargs.hub_arn else None
+
+    if hub_content_type == HubContentType.MODEL_REFERENCE:
+        kwargs.model_reference_arn = construct_hub_model_reference_arn_from_inputs(
+            hub_arn=kwargs.hub_arn, model_name=kwargs.model_id, version=kwargs.model_version
+        )
+    else:
+        kwargs.model_reference_arn = None
+    return kwargs
+
 
 def _add_model_uri_to_kwargs(kwargs: JumpStartEstimatorInitKwargs) -> JumpStartEstimatorInitKwargs:
     """Sets model uri in kwargs based on default or override, returns full kwargs."""

src/sagemaker/jumpstart/factory/model.py

@@ -268,6 +268,7 @@ def _add_model_reference_arn_to_kwargs(
     kwargs: JumpStartModelInitKwargs,
 ) -> JumpStartModelInitKwargs:
     """Sets Model Reference ARN if the hub content type is Model Reference, returns full kwargs."""
+    
     hub_content_type = verify_model_region_and_return_specs(
         model_id=kwargs.model_id,
         version=kwargs.model_version,

src/sagemaker/jumpstart/types.py

@@ -2026,6 +2026,8 @@ class JumpStartEstimatorInitKwargs(JumpStartKwargs):
         "enable_infra_check",
         "enable_remote_debug",
         "enable_session_tag_chaining",
+        "hub_content_type",
+        "model_reference_arn",
     ]
 
     SERIALIZATION_EXCLUSION_SET = {
@@ -2036,6 +2038,7 @@ class JumpStartEstimatorInitKwargs(JumpStartKwargs):
         "model_version",
         "hub_arn",
         "model_type",
+        "hub_content_type",
     }
 
     def __init__(