change: improved jumpstart tagging (#4231)

Author: evakravi

src/sagemaker/estimator.py

@@ -71,7 +71,7 @@
 from sagemaker.utils import instance_supports_kms
 from sagemaker.job import _Job
 from sagemaker.jumpstart.utils import (
-    add_jumpstart_tags,
+    add_jumpstart_uri_tags,
     get_jumpstart_base_name_if_jumpstart_model,
     update_inference_tags_with_jumpstart_training_tags,
 )
@@ -577,9 +577,7 @@ def __init__(
         self.entry_point = entry_point
         self.dependencies = dependencies or []
         self.uploaded_code: Optional[UploadedCode] = None
-        self.tags = add_jumpstart_tags(
-            tags=tags, training_model_uri=self.model_uri, training_script_uri=self.source_dir
-        )
+
         if self.instance_type in ("local", "local_gpu"):
             if self.instance_type == "local_gpu" and self.instance_count > 1:
                 raise RuntimeError("Distributed Training in Local GPU is not supported")
@@ -592,6 +590,15 @@ def __init__(
         else:
             self.sagemaker_session = sagemaker_session or Session()
 
+        self.tags = (
+            add_jumpstart_uri_tags(
+                tags=tags, training_model_uri=self.model_uri, training_script_uri=self.source_dir
+            )
+            if getattr(self.sagemaker_session, "settings", None) is not None
+            and self.sagemaker_session.settings.include_jumpstart_tags
+            else tags
+        )
+
         self.base_job_name = base_job_name
         self._current_job_name = None
         if (

src/sagemaker/jumpstart/constants.py

@@ -154,6 +154,10 @@
     if region.gated_content_bucket is not None
 }
 
+JUMPSTART_GATED_AND_PUBLIC_BUCKET_NAME_SET = JUMPSTART_BUCKET_NAME_SET.union(
+    JUMPSTART_GATED_BUCKET_NAME_SET
+)
+
 JUMPSTART_DEFAULT_REGION_NAME = boto3.session.Session().region_name or "us-west-2"
 
 JUMPSTART_DEFAULT_MANIFEST_FILE_S3_KEY = "models_manifest.json"

src/sagemaker/jumpstart/enums.py

@@ -76,6 +76,9 @@ class JumpStartTag(str, Enum):
     TRAINING_MODEL_URI = "aws-jumpstart-training-model-uri"
     TRAINING_SCRIPT_URI = "aws-jumpstart-training-script-uri"
 
+    MODEL_ID = "sagemaker-sdk:jumpstart-model-id"
+    MODEL_VERSION = "sagemaker-sdk:jumpstart-model-version"
+
 
 class SerializerType(str, Enum):
     """Enum class for serializers associated with JumpStart models."""

src/sagemaker/jumpstart/factory/estimator.py

@@ -60,8 +60,10 @@
     JumpStartModelInitKwargs,
 )
 from sagemaker.jumpstart.utils import (
+    add_jumpstart_model_id_version_tags,
     update_dict_if_key_not_present,
     resolve_estimator_sagemaker_config_field,
+    verify_model_region_and_return_specs,
 )
 
 
@@ -196,6 +198,7 @@ def get_init_kwargs(
     estimator_init_kwargs = _add_estimator_extra_kwargs(estimator_init_kwargs)
     estimator_init_kwargs = _add_role_to_kwargs(estimator_init_kwargs)
     estimator_init_kwargs = _add_env_to_kwargs(estimator_init_kwargs)
+    estimator_init_kwargs = _add_tags_to_kwargs(estimator_init_kwargs)
 
     return estimator_init_kwargs
 
@@ -439,6 +442,26 @@ def _add_instance_type_and_count_to_kwargs(
     return kwargs
 
 
+def _add_tags_to_kwargs(kwargs: JumpStartEstimatorInitKwargs) -> JumpStartEstimatorInitKwargs:
+    """Sets tags in kwargs based on default or override, returns full kwargs."""
+
+    full_model_version = verify_model_region_and_return_specs(
+        model_id=kwargs.model_id,
+        version=kwargs.model_version,
+        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,
+    ).version
+
+    if kwargs.sagemaker_session.settings.include_jumpstart_tags:
+        kwargs.tags = add_jumpstart_model_id_version_tags(
+            kwargs.tags, kwargs.model_id, full_model_version
+        )
+    return kwargs
+
+
 def _add_image_uri_to_kwargs(kwargs: JumpStartEstimatorInitKwargs) -> JumpStartEstimatorInitKwargs:
     """Sets image uri in kwargs based on default or override, returns full kwargs."""
 

src/sagemaker/jumpstart/factory/model.py

@@ -44,6 +44,7 @@
     JumpStartModelRegisterKwargs,
 )
 from sagemaker.jumpstart.utils import (
+    add_jumpstart_model_id_version_tags,
     update_dict_if_key_not_present,
     resolve_model_sagemaker_config_field,
     verify_model_region_and_return_specs,
@@ -422,6 +423,27 @@ def _add_model_name_to_kwargs(
     return kwargs
 
 
+def _add_tags_to_kwargs(kwargs: JumpStartModelDeployKwargs) -> Dict[str, Any]:
+    """Sets tags based on default or override, returns full kwargs."""
+
+    full_model_version = verify_model_region_and_return_specs(
+        model_id=kwargs.model_id,
+        version=kwargs.model_version,
+        scope=JumpStartScriptScope.INFERENCE,
+        region=kwargs.region,
+        tolerate_vulnerable_model=kwargs.tolerate_vulnerable_model,
+        tolerate_deprecated_model=kwargs.tolerate_deprecated_model,
+        sagemaker_session=kwargs.sagemaker_session,
+    ).version
+
+    if kwargs.sagemaker_session.settings.include_jumpstart_tags:
+        kwargs.tags = add_jumpstart_model_id_version_tags(
+            kwargs.tags, kwargs.model_id, full_model_version
+        )
+
+    return kwargs
+
+
 def _add_deploy_extra_kwargs(kwargs: JumpStartModelInitKwargs) -> Dict[str, Any]:
     """Sets extra kwargs based on default or override, returns full kwargs."""
 
@@ -509,6 +531,8 @@ def get_deploy_kwargs(
 
     deploy_kwargs = _add_deploy_extra_kwargs(kwargs=deploy_kwargs)
 
+    deploy_kwargs = _add_tags_to_kwargs(kwargs=deploy_kwargs)
+
     return deploy_kwargs
 
 

src/sagemaker/jumpstart/model.py

@@ -552,6 +552,7 @@ def deploy(
             container_startup_health_check_timeout=container_startup_health_check_timeout,
             inference_recommendation_id=inference_recommendation_id,
             explainer_config=explainer_config,
+            sagemaker_session=self.sagemaker_session,
         )
 
         predictor = super(JumpStartModel, self).deploy(**deploy_kwargs.to_kwargs_dict())

src/sagemaker/jumpstart/utils.py

@@ -250,7 +250,7 @@ def is_jumpstart_model_uri(uri: Optional[str]) -> bool:
     if urlparse(uri).scheme == "s3":
         bucket, _ = parse_s3_url(uri)
 
-    return bucket in constants.JUMPSTART_BUCKET_NAME_SET
+    return bucket in constants.JUMPSTART_GATED_AND_PUBLIC_BUCKET_NAME_SET
 
 
 def tag_key_in_array(tag_key: str, tag_array: List[Dict[str, str]]) -> bool:
@@ -287,7 +287,10 @@ def get_tag_value(tag_key: str, tag_array: List[Dict[str, str]]) -> str:
 
 
 def add_single_jumpstart_tag(
-    uri: str, tag_key: enums.JumpStartTag, curr_tags: Optional[List[Dict[str, str]]]
+    tag_value: str,
+    tag_key: enums.JumpStartTag,
+    curr_tags: Optional[List[Dict[str, str]]],
+    is_uri=False,
 ) -> Optional[List]:
     """Adds ``tag_key`` to ``curr_tags`` if ``uri`` corresponds to a JumpStart model.
 
@@ -296,17 +299,28 @@ def add_single_jumpstart_tag(
         tag_key (enums.JumpStartTag): Custom tag to apply to current tags if the URI
             corresponds to a JumpStart model.
         curr_tags (Optional[List]): Current tags associated with ``Estimator`` or ``Model``.
+        is_uri (boolean): Set to True to indicate a s3 uri is to be tagged. Set to False to indicate
+            tags for JumpStart model id / version are being added. (Default: False).
     """
-    if is_jumpstart_model_uri(uri):
+    if not is_uri or is_jumpstart_model_uri(tag_value):
         if curr_tags is None:
             curr_tags = []
         if not tag_key_in_array(tag_key, curr_tags):
-            curr_tags.append(
-                {
-                    "Key": tag_key,
-                    "Value": uri,
-                }
+            skip_adding_tag = (
+                (
+                    tag_key_in_array(enums.JumpStartTag.MODEL_ID, curr_tags)
+                    or tag_key_in_array(enums.JumpStartTag.MODEL_VERSION, curr_tags)
+                )
+                if is_uri
+                else False
             )
+            if not skip_adding_tag:
+                curr_tags.append(
+                    {
+                        "Key": tag_key,
+                        "Value": tag_value,
+                    }
+                )
     return curr_tags
 
 
@@ -326,14 +340,37 @@ def get_jumpstart_base_name_if_jumpstart_model(
     return None
 
 
-def add_jumpstart_tags(
+def add_jumpstart_model_id_version_tags(
+    tags: Optional[List[Dict[str, str]]],
+    model_id: str,
+    model_version: str,
+) -> List[Dict[str, str]]:
+    """Add custom model ID and version tags to JumpStart related resources."""
+    if model_id is None or model_version is None:
+        return tags
+    tags = add_single_jumpstart_tag(
+        model_id,
+        enums.JumpStartTag.MODEL_ID,
+        tags,
+        is_uri=False,
+    )
+    tags = add_single_jumpstart_tag(
+        model_version,
+        enums.JumpStartTag.MODEL_VERSION,
+        tags,
+        is_uri=False,
+    )
+    return tags
+
+
+def add_jumpstart_uri_tags(
     tags: Optional[List[Dict[str, str]]] = None,
     inference_model_uri: Optional[Union[str, dict]] = None,
     inference_script_uri: Optional[str] = None,
     training_model_uri: Optional[str] = None,
     training_script_uri: Optional[str] = None,
 ) -> Optional[List[Dict[str, str]]]:
-    """Add custom tags to JumpStart models, return the updated tags.
+    """Add custom uri tags to JumpStart models, return the updated tags.
 
     No-op if this is not a JumpStart model related resource.
 
@@ -362,31 +399,43 @@ def add_jumpstart_tags(
             logging.warning(warn_msg, "inference_model_uri")
         else:
             tags = add_single_jumpstart_tag(
-                inference_model_uri, enums.JumpStartTag.INFERENCE_MODEL_URI, tags
+                inference_model_uri,
+                enums.JumpStartTag.INFERENCE_MODEL_URI,
+                tags,
+                is_uri=True,
             )
 
     if inference_script_uri:
         if is_pipeline_variable(inference_script_uri):
             logging.warning(warn_msg, "inference_script_uri")
         else:
             tags = add_single_jumpstart_tag(
-                inference_script_uri, enums.JumpStartTag.INFERENCE_SCRIPT_URI, tags
+                inference_script_uri,
+                enums.JumpStartTag.INFERENCE_SCRIPT_URI,
+                tags,
+                is_uri=True,
             )
 
     if training_model_uri:
         if is_pipeline_variable(training_model_uri):
             logging.warning(warn_msg, "training_model_uri")
         else:
             tags = add_single_jumpstart_tag(
-                training_model_uri, enums.JumpStartTag.TRAINING_MODEL_URI, tags
+                training_model_uri,
+                enums.JumpStartTag.TRAINING_MODEL_URI,
+                tags,
+                is_uri=True,
             )
 
     if training_script_uri:
         if is_pipeline_variable(training_script_uri):
             logging.warning(warn_msg, "training_script_uri")
         else:
             tags = add_single_jumpstart_tag(
-                training_script_uri, enums.JumpStartTag.TRAINING_SCRIPT_URI, tags
+                training_script_uri,
+                enums.JumpStartTag.TRAINING_SCRIPT_URI,
+                tags,
+                is_uri=True,
             )
 
     return tags

src/sagemaker/model.py

@@ -54,7 +54,7 @@
 from sagemaker.serverless import ServerlessInferenceConfig
 from sagemaker.transformer import Transformer
 from sagemaker.jumpstart.utils import (
-    add_jumpstart_tags,
+    add_jumpstart_uri_tags,
     get_jumpstart_base_name_if_jumpstart_model,
 )
 from sagemaker.utils import (
@@ -1375,13 +1375,17 @@ def deploy(
             sagemaker_session=self.sagemaker_session,
         )
 
-        tags = add_jumpstart_tags(
-            tags=tags,
-            inference_model_uri=self.model_data
-            if isinstance(self.model_data, (str, dict))
-            else None,
-            inference_script_uri=self.source_dir,
-        )
+        if (
+            getattr(self.sagemaker_session, "settings", None) is not None
+            and self.sagemaker_session.settings.include_jumpstart_tags
+        ):
+            tags = add_jumpstart_uri_tags(
+                tags=tags,
+                inference_model_uri=self.model_data
+                if isinstance(self.model_data, (str, dict))
+                else None,
+                inference_script_uri=self.source_dir,
+            )
 
         if self.role is None:
             raise ValueError("Role can not be null for deploying a model")

src/sagemaker/session_settings.py

@@ -22,6 +22,7 @@ def __init__(
         self,
         encrypt_repacked_artifacts=True,
         local_download_dir=None,
+        include_jumpstart_tags=True,
     ) -> None:
         """Initialize the ``SessionSettings`` of a SageMaker ``Session``.
 
@@ -31,9 +32,12 @@ def __init__(
                 is not provided (Default: True).
             local_download_dir (str): Optional. A path specifying the local directory
                 for downloading artifacts. (Default: None).
+            include_jumpstart_tags (bool): Optional. By default, if a JumpStart model is identified,
+                it will receive special tags describing its properties.
         """
         self._encrypt_repacked_artifacts = encrypt_repacked_artifacts
         self._local_download_dir = local_download_dir
+        self._include_jumpstart_tags = include_jumpstart_tags
 
     @property
     def encrypt_repacked_artifacts(self) -> bool:
@@ -44,3 +48,8 @@ def encrypt_repacked_artifacts(self) -> bool:
     def local_download_dir(self) -> str:
         """Return path specifying the local directory for downloading artifacts."""
         return self._local_download_dir
+
+    @property
+    def include_jumpstart_tags(self) -> bool:
+        """Return True if JumpStart tags should be attached to models with JumpStart artifacts."""
+        return self._include_jumpstart_tags

src/sagemaker/tuner.py

@@ -34,7 +34,7 @@
 from sagemaker.inputs import TrainingInput, FileSystemInput
 from sagemaker.job import _Job
 from sagemaker.jumpstart.utils import (
-    add_jumpstart_tags,
+    add_jumpstart_uri_tags,
     get_jumpstart_base_name_if_jumpstart_model,
 )
 from sagemaker.parameter import (
@@ -816,11 +816,12 @@ def _prepare_tags_for_tuning(self):
             if tag not in self.tags:
                 self.tags.append(tag)
 
-        self.tags = add_jumpstart_tags(
-            tags=self.tags,
-            training_script_uri=getattr(estimator, "source_dir", None),
-            training_model_uri=self._get_model_uri(estimator),
-        )
+        if self.sagemaker_session.settings.include_jumpstart_tags:
+            self.tags = add_jumpstart_uri_tags(
+                tags=self.tags,
+                training_script_uri=getattr(estimator, "source_dir", None),
+                training_model_uri=self._get_model_uri(estimator),
+            )
 
     def _prepare_job_name_for_tuning(self, job_name=None):
         """Set current job name before starting tuning."""