Feat: tagging jumpstart models (#2860)

Author: evakravi

src/sagemaker/estimator.py

@@ -48,6 +48,10 @@
 )
 from sagemaker.inputs import TrainingInput
 from sagemaker.job import _Job
+from sagemaker.jumpstart.utils import (
+    add_jumpstart_tags,
+    update_inference_tags_with_jumpstart_training_tags,
+)
 from sagemaker.local import LocalSession
 from sagemaker.model import (
     CONTAINER_LOG_LEVEL_PARAM_NAME,
@@ -442,7 +446,6 @@ def __init__(
         self.volume_kms_key = volume_kms_key
         self.max_run = max_run
         self.input_mode = input_mode
-        self.tags = tags
         self.metric_definitions = metric_definitions
         self.model_uri = model_uri
         self.model_channel_name = model_channel_name
@@ -456,7 +459,9 @@ def __init__(
         self.entry_point = entry_point
         self.dependencies = dependencies
         self.uploaded_code = 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")
@@ -1203,6 +1208,10 @@ def deploy(
 
         model.name = model_name
 
+        tags = update_inference_tags_with_jumpstart_training_tags(
+            inference_tags=tags, training_tags=self.tags
+        )
+
         return model.deploy(
             instance_type=instance_type,
             initial_instance_count=initial_instance_count,

src/sagemaker/jumpstart/constants.py

@@ -112,6 +112,8 @@
 }
 JUMPSTART_REGION_NAME_SET = {region.region_name for region in JUMPSTART_LAUNCHED_REGIONS}
 
+JUMPSTART_BUCKET_NAME_SET = {region.content_bucket for region in JUMPSTART_LAUNCHED_REGIONS}
+
 JUMPSTART_DEFAULT_REGION_NAME = boto3.session.Session().region_name or "us-west-2"
 
 JUMPSTART_DEFAULT_MANIFEST_FILE_S3_KEY = "models_manifest.json"
@@ -156,3 +158,12 @@ class VariableScope(str, Enum):
 
     CONTAINER = "container"
     ALGORITHM = "algorithm"
+
+
+class JumpStartTag(str, Enum):
+    """Enum class for tag keys to apply to JumpStart models."""
+
+    INFERENCE_MODEL_URI = "aws-jumpstart-inference-model-uri"
+    INFERENCE_SCRIPT_URI = "aws-jumpstart-inference-script-uri"
+    TRAINING_MODEL_URI = "aws-jumpstart-training-model-uri"
+    TRAINING_SCRIPT_URI = "aws-jumpstart-training-script-uri"

src/sagemaker/jumpstart/utils.py

@@ -14,10 +14,12 @@
 from __future__ import absolute_import
 import logging
 from typing import Dict, List, Optional
+from urllib.parse import urlparse
 from packaging.version import Version
 import sagemaker
 from sagemaker.jumpstart import constants
 from sagemaker.jumpstart import accessors
+from sagemaker.s3 import parse_s3_url
 from sagemaker.jumpstart.exceptions import (
     DeprecatedJumpStartModelError,
     VulnerableJumpStartModelError,
@@ -150,6 +152,145 @@ def is_jumpstart_model_input(model_id: Optional[str], version: Optional[str]) ->
     return False
 
 
+def is_jumpstart_model_uri(uri: Optional[str]) -> bool:
+    """Returns True if URI corresponds to a JumpStart-hosted model.
+
+    Args:
+        uri (Optional[str]): uri for inference/training job.
+    """
+
+    bucket = None
+    if urlparse(uri).scheme == "s3":
+        bucket, _ = parse_s3_url(uri)
+
+    return bucket in constants.JUMPSTART_BUCKET_NAME_SET
+
+
+def tag_key_in_array(tag_key: str, tag_array: List[Dict[str, str]]) -> bool:
+    """Returns True if ``tag_key`` is in the ``tag_array``.
+
+    Args:
+        tag_key (str): the tag key to check if it's already in the ``tag_array``.
+        tag_array (List[Dict[str, str]]): array of tags to check for ``tag_key``.
+    """
+    for tag in tag_array:
+        if tag_key == tag["Key"]:
+            return True
+    return False
+
+
+def get_tag_value(tag_key: str, tag_array: List[Dict[str, str]]) -> str:
+    """Return the value of a tag whose key matches the given ``tag_key``.
+
+    Args:
+        tag_key (str): AWS tag for which to search.
+        tag_array (List[Dict[str, str]]): List of AWS tags, each formatted as dicts.
+
+    Raises:
+        KeyError: If the number of matches for the ``tag_key`` is not equal to 1.
+    """
+    tag_values = [tag["Value"] for tag in tag_array if tag_key == tag["Key"]]
+    if len(tag_values) != 1:
+        raise KeyError(
+            f"Cannot get value of tag for tag key '{tag_key}' -- found {len(tag_values)} "
+            f"number of matches in the tag list."
+        )
+
+    return tag_values[0]
+
+
+def add_single_jumpstart_tag(
+    uri: str, tag_key: constants.JumpStartTag, curr_tags: Optional[List[Dict[str, str]]]
+) -> Optional[List]:
+    """Adds ``tag_key`` to ``curr_tags`` if ``uri`` corresponds to a JumpStart model.
+
+    Args:
+        uri (str): URI which may correspond to a JumpStart model.
+        tag_key (constants.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``.
+    """
+    if is_jumpstart_model_uri(uri):
+        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,
+                }
+            )
+    return curr_tags
+
+
+def add_jumpstart_tags(
+    tags: Optional[List[Dict[str, str]]] = None,
+    inference_model_uri: Optional[str] = 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.
+
+    No-op if this is not a JumpStart model related resource.
+
+    Args:
+        tags (Optional[List[Dict[str,str]]): Current tags for JumpStart inference
+            or training job. (Default: None).
+        inference_model_uri (Optional[str]): S3 URI for inference model artifact.
+            (Default: None).
+        inference_script_uri (Optional[str]): S3 URI for inference script tarball.
+            (Default: None).
+        training_model_uri (Optional[str]): S3 URI for training model artifact.
+            (Default: None).
+        training_script_uri (Optional[str]): S3 URI for training script tarball.
+            (Default: None).
+    """
+
+    if inference_model_uri:
+        tags = add_single_jumpstart_tag(
+            inference_model_uri, constants.JumpStartTag.INFERENCE_MODEL_URI, tags
+        )
+
+    if inference_script_uri:
+        tags = add_single_jumpstart_tag(
+            inference_script_uri, constants.JumpStartTag.INFERENCE_SCRIPT_URI, tags
+        )
+
+    if training_model_uri:
+        tags = add_single_jumpstart_tag(
+            training_model_uri, constants.JumpStartTag.TRAINING_MODEL_URI, tags
+        )
+
+    if training_script_uri:
+        tags = add_single_jumpstart_tag(
+            training_script_uri, constants.JumpStartTag.TRAINING_SCRIPT_URI, tags
+        )
+
+    return tags
+
+
+def update_inference_tags_with_jumpstart_training_tags(
+    inference_tags: Optional[List[Dict[str, str]]], training_tags: Optional[List[Dict[str, str]]]
+) -> Optional[List[Dict[str, str]]]:
+    """Updates the tags for the ``sagemaker.model.Model.deploy`` command with any JumpStart tags.
+
+    Args:
+        inference_tags (Optional[List[Dict[str, str]]]): Custom tags to appy to inference job.
+        training_tags (Optional[List[Dict[str, str]]]): Tags from training job.
+    """
+    if training_tags:
+        for tag_key in constants.JumpStartTag:
+            if tag_key_in_array(tag_key, training_tags):
+                tag_value = get_tag_value(tag_key, training_tags)
+                if inference_tags is None:
+                    inference_tags = []
+                if not tag_key_in_array(tag_key, inference_tags):
+                    inference_tags.append({"Key": tag_key, "Value": tag_value})
+
+    return inference_tags
+
+
 def verify_model_region_and_return_specs(
     model_id: Optional[str],
     version: Optional[str],

src/sagemaker/model.py

@@ -34,6 +34,7 @@
 from sagemaker.deprecations import removed_kwargs
 from sagemaker.predictor import PredictorBase
 from sagemaker.transformer import Transformer
+from sagemaker.jumpstart.utils import add_jumpstart_tags
 
 LOGGER = logging.getLogger("sagemaker")
 
@@ -985,6 +986,10 @@ def deploy(
         removed_kwargs("update_endpoint", kwargs)
         self._init_sagemaker_session_if_does_not_exist(instance_type)
 
+        tags = add_jumpstart_tags(
+            tags=tags, inference_model_uri=self.model_data, inference_script_uri=self.source_dir
+        )
+
         if self.role is None:
             raise ValueError("Role can not be null for deploying a model")