feat: tagging jumpstart models

evakravi · evakravi · commit 0ee60659afce · 2022-01-19T15:04:24.000Z
diff --git a/src/sagemaker/jumpstart/constants.py b/src/sagemaker/jumpstart/constants.py
@@ -12,18 +12,115 @@
 # language governing permissions and limitations under the License.
 """This module stores constants related to SageMaker JumpStart."""
 from __future__ import absolute_import
+from enum import Enum
 from typing import Set
 import boto3
 from sagemaker.jumpstart.types import JumpStartLaunchedRegionInfo
 
 
-JUMPSTART_LAUNCHED_REGIONS: Set[JumpStartLaunchedRegionInfo] = set()
+JUMPSTART_LAUNCHED_REGIONS: Set[JumpStartLaunchedRegionInfo] = set(
+    [
+        JumpStartLaunchedRegionInfo(
+            region_name="us-west-2",
+            content_bucket="jumpstart-cache-prod-us-west-2",
+        ),
+        JumpStartLaunchedRegionInfo(
+            region_name="us-east-1",
+            content_bucket="jumpstart-cache-prod-us-east-1",
+        ),
+        JumpStartLaunchedRegionInfo(
+            region_name="us-east-2",
+            content_bucket="jumpstart-cache-prod-us-east-2",
+        ),
+        JumpStartLaunchedRegionInfo(
+            region_name="eu-west-1",
+            content_bucket="jumpstart-cache-prod-eu-west-1",
+        ),
+        JumpStartLaunchedRegionInfo(
+            region_name="eu-central-1",
+            content_bucket="jumpstart-cache-prod-eu-central-1",
+        ),
+        JumpStartLaunchedRegionInfo(
+            region_name="eu-north-1",
+            content_bucket="jumpstart-cache-prod-eu-north-1",
+        ),
+        JumpStartLaunchedRegionInfo(
+            region_name="me-south-1",
+            content_bucket="jumpstart-cache-prod-me-south-1",
+        ),
+        JumpStartLaunchedRegionInfo(
+            region_name="ap-south-1",
+            content_bucket="jumpstart-cache-prod-ap-south-1",
+        ),
+        JumpStartLaunchedRegionInfo(
+            region_name="eu-west-3",
+            content_bucket="jumpstart-cache-prod-eu-west-3",
+        ),
+        JumpStartLaunchedRegionInfo(
+            region_name="af-south-1",
+            content_bucket="jumpstart-cache-prod-af-south-1",
+        ),
+        JumpStartLaunchedRegionInfo(
+            region_name="sa-east-1",
+            content_bucket="jumpstart-cache-prod-sa-east-1",
+        ),
+        JumpStartLaunchedRegionInfo(
+            region_name="ap-east-1",
+            content_bucket="jumpstart-cache-prod-ap-east-1",
+        ),
+        JumpStartLaunchedRegionInfo(
+            region_name="ap-northeast-2",
+            content_bucket="jumpstart-cache-prod-ap-northeast-2",
+        ),
+        JumpStartLaunchedRegionInfo(
+            region_name="eu-west-2",
+            content_bucket="jumpstart-cache-prod-eu-west-2",
+        ),
+        JumpStartLaunchedRegionInfo(
+            region_name="eu-south-1",
+            content_bucket="jumpstart-cache-prod-eu-south-1",
+        ),
+        JumpStartLaunchedRegionInfo(
+            region_name="ap-northeast-1",
+            content_bucket="jumpstart-cache-prod-ap-northeast-1",
+        ),
+        JumpStartLaunchedRegionInfo(
+            region_name="us-west-1",
+            content_bucket="jumpstart-cache-prod-us-west-1",
+        ),
+        JumpStartLaunchedRegionInfo(
+            region_name="ap-southeast-1",
+            content_bucket="jumpstart-cache-prod-ap-southeast-1",
+        ),
+        JumpStartLaunchedRegionInfo(
+            region_name="ap-southeast-2",
+            content_bucket="jumpstart-cache-prod-ap-southeast-2",
+        ),
+        JumpStartLaunchedRegionInfo(
+            region_name="ca-central-1",
+            content_bucket="jumpstart-cache-prod-ca-central-1",
+        ),
+        JumpStartLaunchedRegionInfo(
+            region_name="cn-north-1",
+            content_bucket="jumpstart-cache-prod-cn-north-1",
+        ),
+    ]
+)
 
 JUMPSTART_REGION_NAME_TO_LAUNCHED_REGION_DICT = {
     region.region_name: region for region in JUMPSTART_LAUNCHED_REGIONS
 }
 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
 
 JUMPSTART_DEFAULT_MANIFEST_FILE_S3_KEY = "models_manifest.json"
+
+
+class JumpStartTag(str, Enum):
+    """Enum class for tags to apply to JumpStart models."""
+
+    INFERENCE_MODEL_URI = "jumpstart-inference-model-uri"
+    INFERENCE_SCRIPT_URI = "jumpstart-inference-script-uri"
diff --git a/src/sagemaker/jumpstart/utils.py b/src/sagemaker/jumpstart/utils.py
@@ -12,11 +12,13 @@
 # language governing permissions and limitations under the License.
 """This module contains utilities related to SageMaker JumpStart."""
 from __future__ import absolute_import
-from typing import Dict, List
+from typing import Dict, List, Optional
+from urllib.parse import urlparse
 import semantic_version
 import sagemaker
 from sagemaker.jumpstart import constants
 from sagemaker.jumpstart.types import JumpStartModelHeader, JumpStartVersionedModelId
+from sagemaker.s3 import parse_s3_url
 
 
 class SageMakerSettings(object):
@@ -128,3 +130,53 @@ def parse_sagemaker_version() -> str:
     semantic_version.Version(parsed_version)
 
     return parsed_version
+
+
+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 add_jumpstart_tags(
+    tags: Optional[List[Dict[str, str]]],
+    inference_model_uri: Optional[str],
+    inference_script_uri: Optional[str],
+) -> List[Dict[str, str]]:
+    """Adds tags for JumpStart models. Returns original tags for non-JumpStart
+    models.
+
+    Args:
+        tags (Optional[List[Dict[str,str]]): Current tags for JumpStart inference
+        or training job.
+        inference_model_uri (Optional[str]): S3 URI for inference model artifact.
+        inference_script_uri (Optional[str]): S3 URI for inference script tarball.
+    """
+
+    if is_jumpstart_model_uri(inference_model_uri):
+        if tags is None:
+            tags = []
+        tags.append(
+            {
+                constants.JumpStartTag.INFERENCE_MODEL_URI.value: inference_model_uri,
+            }
+        )
+
+    if is_jumpstart_model_uri(inference_script_uri):
+        if tags is None:
+            tags = []
+        tags.append(
+            {
+                constants.JumpStartTag.INFERENCE_SCRIPT_URI.value: inference_script_uri,
+            }
+        )
+
+    return tags
diff --git a/src/sagemaker/model.py b/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")
 
@@ -982,6 +983,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")
 
diff --git a/tests/unit/sagemaker/jumpstart/test_utils.py b/tests/unit/sagemaker/jumpstart/test_utils.py
@@ -13,11 +13,20 @@
 from __future__ import absolute_import
 from mock.mock import Mock, patch
 import pytest
+import random
 from sagemaker.jumpstart import utils
-from sagemaker.jumpstart.constants import JUMPSTART_REGION_NAME_SET
+from sagemaker.jumpstart.constants import (
+    JUMPSTART_BUCKET_NAME_SET,
+    JUMPSTART_REGION_NAME_SET,
+    JumpStartTag,
+)
 from sagemaker.jumpstart.types import JumpStartModelHeader, JumpStartVersionedModelId
 
 
+def random_jumpstart_s3_uri(key):
+    return f"s3://{random.choice(list(JUMPSTART_BUCKET_NAME_SET))}/{key}"
+
+
 def test_get_jumpstart_content_bucket():
     bad_region = "bad_region"
     assert bad_region not in JUMPSTART_REGION_NAME_SET
@@ -112,3 +121,168 @@ def test_get_sagemaker_version(patched_parse_sm_version: Mock):
     utils.get_sagemaker_version()
     utils.get_sagemaker_version()
     assert patched_parse_sm_version.called_only_once()
+
+
+def test_is_jumpstart_model_uri():
+
+    assert not utils.is_jumpstart_model_uri("fdsfdsf")
+    assert not utils.is_jumpstart_model_uri("s3://not-jumpstart-bucket/sdfsdfds")
+    assert not utils.is_jumpstart_model_uri("some/actual/localfile")
+
+    assert utils.is_jumpstart_model_uri(
+        random_jumpstart_s3_uri("source_directory_tarballs/sourcedir.tar.gz")
+    )
+    assert utils.is_jumpstart_model_uri(random_jumpstart_s3_uri("random_key"))
+
+
+def test_add_jumpstart_tags():
+    tags = None
+    inference_model_uri = "dfsdfsd"
+    inference_script_uri = "dfsdfs"
+    assert (
+        utils.add_jumpstart_tags(
+            tags=tags,
+            inference_model_uri=inference_model_uri,
+            inference_script_uri=inference_script_uri,
+        )
+        is None
+    )
+
+    tags = []
+    inference_model_uri = "dfsdfsd"
+    inference_script_uri = "dfsdfs"
+    assert (
+        utils.add_jumpstart_tags(
+            tags=tags,
+            inference_model_uri=inference_model_uri,
+            inference_script_uri=inference_script_uri,
+        )
+        == []
+    )
+
+    tags = [{"some": "tag"}]
+    inference_model_uri = "dfsdfsd"
+    inference_script_uri = "dfsdfs"
+    assert (
+        utils.add_jumpstart_tags(
+            tags=tags,
+            inference_model_uri=inference_model_uri,
+            inference_script_uri=inference_script_uri,
+        )
+        == [{"some": "tag"}]
+    )
+
+    tags = None
+    inference_model_uri = random_jumpstart_s3_uri("random_key")
+    inference_script_uri = "dfsdfs"
+    assert (
+        utils.add_jumpstart_tags(
+            tags=tags,
+            inference_model_uri=inference_model_uri,
+            inference_script_uri=inference_script_uri,
+        )
+        == [{JumpStartTag.INFERENCE_MODEL_URI.value: inference_model_uri}]
+    )
+
+    tags = []
+    inference_model_uri = random_jumpstart_s3_uri("random_key")
+    inference_script_uri = "dfsdfs"
+    assert (
+        utils.add_jumpstart_tags(
+            tags=tags,
+            inference_model_uri=inference_model_uri,
+            inference_script_uri=inference_script_uri,
+        )
+        == [{JumpStartTag.INFERENCE_MODEL_URI.value: inference_model_uri}]
+    )
+
+    tags = [{"some": "tag"}]
+    inference_model_uri = random_jumpstart_s3_uri("random_key")
+    inference_script_uri = "dfsdfs"
+    assert utils.add_jumpstart_tags(
+        tags=tags,
+        inference_model_uri=inference_model_uri,
+        inference_script_uri=inference_script_uri,
+    ) == [
+        {"some": "tag"},
+        {JumpStartTag.INFERENCE_MODEL_URI.value: inference_model_uri},
+    ]
+
+    tags = None
+    inference_script_uri = random_jumpstart_s3_uri("random_key")
+    inference_model_uri = "dfsdfs"
+    assert (
+        utils.add_jumpstart_tags(
+            tags=tags,
+            inference_model_uri=inference_model_uri,
+            inference_script_uri=inference_script_uri,
+        )
+        == [{JumpStartTag.INFERENCE_SCRIPT_URI.value: inference_script_uri}]
+    )
+
+    tags = []
+    inference_script_uri = random_jumpstart_s3_uri("random_key")
+    inference_model_uri = "dfsdfs"
+    assert (
+        utils.add_jumpstart_tags(
+            tags=tags,
+            inference_model_uri=inference_model_uri,
+            inference_script_uri=inference_script_uri,
+        )
+        == [{JumpStartTag.INFERENCE_SCRIPT_URI.value: inference_script_uri}]
+    )
+
+    tags = [{"some": "tag"}]
+    inference_script_uri = random_jumpstart_s3_uri("random_key")
+    inference_model_uri = "dfsdfs"
+    assert utils.add_jumpstart_tags(
+        tags=tags,
+        inference_model_uri=inference_model_uri,
+        inference_script_uri=inference_script_uri,
+    ) == [
+        {"some": "tag"},
+        {JumpStartTag.INFERENCE_SCRIPT_URI.value: inference_script_uri},
+    ]
+
+    tags = None
+    inference_script_uri = random_jumpstart_s3_uri("random_key")
+    inference_model_uri = random_jumpstart_s3_uri("random_key")
+    assert utils.add_jumpstart_tags(
+        tags=tags,
+        inference_model_uri=inference_model_uri,
+        inference_script_uri=inference_script_uri,
+    ) == [
+        {
+            JumpStartTag.INFERENCE_MODEL_URI.value: inference_model_uri,
+        },
+        {JumpStartTag.INFERENCE_SCRIPT_URI.value: inference_script_uri},
+    ]
+
+    tags = []
+    inference_script_uri = random_jumpstart_s3_uri("random_key")
+    inference_model_uri = random_jumpstart_s3_uri("random_key")
+    assert utils.add_jumpstart_tags(
+        tags=tags,
+        inference_model_uri=inference_model_uri,
+        inference_script_uri=inference_script_uri,
+    ) == [
+        {
+            JumpStartTag.INFERENCE_MODEL_URI.value: inference_model_uri,
+        },
+        {JumpStartTag.INFERENCE_SCRIPT_URI.value: inference_script_uri},
+    ]
+
+    tags = [{"some": "tag"}]
+    inference_script_uri = random_jumpstart_s3_uri("random_key")
+    inference_model_uri = random_jumpstart_s3_uri("random_key")
+    assert utils.add_jumpstart_tags(
+        tags=tags,
+        inference_model_uri=inference_model_uri,
+        inference_script_uri=inference_script_uri,
+    ) == [
+        {"some": "tag"},
+        {
+            JumpStartTag.INFERENCE_MODEL_URI.value: inference_model_uri,
+        },
+        {JumpStartTag.INFERENCE_SCRIPT_URI.value: inference_script_uri},
+    ]
diff --git a/tests/unit/sagemaker/model/test_model.py b/tests/unit/sagemaker/model/test_model.py
