More tags formatting and add a test

Author: martinRenou

src/sagemaker/algorithm.py

@@ -392,7 +392,7 @@ def transformer(
             if self._is_marketplace():
                 transform_env = None
 
-            tags = tags or self.tags
+            tags = format_tags(tags) or self.tags
         else:
             raise RuntimeError("No finished training job found associated with this estimator")
 

src/sagemaker/apiutils/_base_types.py

@@ -14,6 +14,7 @@
 from __future__ import absolute_import
 
 from sagemaker.apiutils import _boto_functions, _utils
+from sagemaker.utils import format_tags
 
 
 class ApiObject(object):
@@ -194,13 +195,13 @@ def _set_tags(self, resource_arn=None, tags=None):
 
         Args:
             resource_arn (str): The arn of the Record
-            tags (dict): An array of Tag objects that set to Record
+            tags (Optional[Tags]): An array of Tag objects that set to Record
 
         Returns:
             A list of key, value pair objects. i.e. [{"key":"value"}]
         """
         tag_list = self.sagemaker_session.sagemaker_client.add_tags(
-            ResourceArn=resource_arn, Tags=tags
+            ResourceArn=resource_arn, Tags=format_tags(tags)
         )["Tags"]
         return tag_list
 

src/sagemaker/automl/automl.py

@@ -580,7 +580,7 @@ def deploy(
                 be selected on each ``deploy``.
             endpoint_name (str): The name of the endpoint to create (default:
                 None). If not specified, a unique endpoint name will be created.
-            tags (List[dict[str, str]]): The list of tags to attach to this
+            tags (Optional[Tags]): The list of tags to attach to this
                 specific endpoint.
             wait (bool): Whether the call should wait until the deployment of
                 model completes (default: True).
@@ -632,7 +632,7 @@ def deploy(
             deserializer=deserializer,
             endpoint_name=endpoint_name,
             kms_key=model_kms_key,
-            tags=tags,
+            tags=format_tags(tags),
             wait=wait,
             volume_size=volume_size,
             model_data_download_timeout=model_data_download_timeout,

src/sagemaker/base_predictor.py

@@ -53,7 +53,7 @@
     NumpySerializer,
 )
 from sagemaker.session import production_variant, Session
-from sagemaker.utils import name_from_base, stringify_object
+from sagemaker.utils import name_from_base, stringify_object, format_tags
 
 from sagemaker.model_monitor.model_monitoring import DEFAULT_REPOSITORY_NAME
 
@@ -409,7 +409,7 @@ def update_endpoint(
         self.sagemaker_session.create_endpoint_config_from_existing(
             current_endpoint_config_name,
             new_endpoint_config_name,
-            new_tags=tags,
+            new_tags=format_tags(tags),
             new_kms_key=kms_key,
             new_data_capture_config_dict=data_capture_config_dict,
             new_production_variants=production_variants,

src/sagemaker/djl_inference/model.py

@@ -30,7 +30,7 @@
 from sagemaker.s3_utils import s3_path_join
 from sagemaker.serializers import JSONSerializer, BaseSerializer
 from sagemaker.session import Session
-from sagemaker.utils import _tmpdir, _create_or_update_code_dir
+from sagemaker.utils import _tmpdir, _create_or_update_code_dir, format_tags
 from sagemaker.workflow.entities import PipelineVariable
 from sagemaker.estimator import Estimator
 from sagemaker.s3 import S3Uploader
@@ -610,7 +610,7 @@ def deploy(
                 default deserializer is set by the ``predictor_cls``.
             endpoint_name (str): The name of the endpoint to create (default:
                 None). If not specified, a unique endpoint name will be created.
-            tags (List[dict[str, str]]): The list of tags to attach to this
+            tags (Optional[Tags]): The list of tags to attach to this
                 specific endpoint.
             kms_key (str): The ARN of the KMS key that is used to encrypt the
                 data on the storage volume attached to the instance hosting the
@@ -651,7 +651,7 @@ def deploy(
             serializer=serializer,
             deserializer=deserializer,
             endpoint_name=endpoint_name,
-            tags=tags,
+            tags=format_tags(tags),
             kms_key=kms_key,
             wait=wait,
             data_capture_config=data_capture_config,

src/sagemaker/experiments/experiment.py

@@ -20,6 +20,7 @@
 from sagemaker.apiutils import _base_types
 from sagemaker.experiments.trial import _Trial
 from sagemaker.experiments.trial_component import _TrialComponent
+from sagemaker.utils import format_tags
 
 
 class Experiment(_base_types.Record):
@@ -111,7 +112,7 @@ def create(
                 manages interactions with Amazon SageMaker APIs and any other
                 AWS services needed. If not specified, one is created using the
                 default AWS configuration chain.
-            tags (List[Dict[str, str]]): A list of tags to associate with the experiment
+            tags (Optional[Tags]): A list of tags to associate with the experiment
                 (default: None).
 
         Returns:
@@ -122,7 +123,7 @@ def create(
             experiment_name=experiment_name,
             display_name=display_name,
             description=description,
-            tags=tags,
+            tags=format_tags(tags),
             sagemaker_session=sagemaker_session,
         )
 
@@ -149,7 +150,7 @@ def _load_or_create(
                 manages interactions with Amazon SageMaker APIs and any other
                 AWS services needed. If not specified, one is created using the
                 default AWS configuration chain.
-            tags (List[Dict[str, str]]): A list of tags to associate with the experiment
+            tags (Optional[Tags]): A list of tags to associate with the experiment
                 (default: None). This is used only when the given `experiment_name` does not
                 exist and a new experiment has to be created.
 
@@ -161,7 +162,7 @@ def _load_or_create(
                 experiment_name=experiment_name,
                 display_name=display_name,
                 description=description,
-                tags=tags,
+                tags=format_tags(tags),
                 sagemaker_session=sagemaker_session,
             )
         except ClientError as ce:

src/sagemaker/experiments/trial.py

@@ -18,6 +18,7 @@
 from sagemaker.apiutils import _base_types
 from sagemaker.experiments import _api_types
 from sagemaker.experiments.trial_component import _TrialComponent
+from sagemaker.utils import format_tags
 
 
 class _Trial(_base_types.Record):
@@ -101,7 +102,7 @@ def create(
             trial_name: (str): Name of the Trial.
             display_name (str): Name of the trial that will appear in UI,
                 such as SageMaker Studio (default: None).
-            tags (List[dict]): A list of tags to associate with the trial (default: None).
+            tags (Optional[Tags]): A list of tags to associate with the trial (default: None).
             sagemaker_session (sagemaker.session.Session): Session object which
                 manages interactions with Amazon SageMaker APIs and any other
                 AWS services needed. If not specified, one is created using the
@@ -115,7 +116,7 @@ def create(
             trial_name=trial_name,
             experiment_name=experiment_name,
             display_name=display_name,
-            tags=tags,
+            tags=format_tags(tags),
             sagemaker_session=sagemaker_session,
         )
         return trial
@@ -259,7 +260,7 @@ def _load_or_create(
             display_name (str): Name of the trial that will appear in UI,
                 such as SageMaker Studio (default: None). This is used only when the given
                 `trial_name` does not exist and a new trial has to be created.
-            tags (List[dict]): A list of tags to associate with the trial (default: None).
+            tags (Optional[Tags]): A list of tags to associate with the trial (default: None).
                 This is used only when the given `trial_name` does not exist and
                 a new trial has to be created.
             sagemaker_session (sagemaker.session.Session): Session object which
@@ -275,7 +276,7 @@ def _load_or_create(
                 experiment_name=experiment_name,
                 trial_name=trial_name,
                 display_name=display_name,
-                tags=tags,
+                tags=format_tags(tags),
                 sagemaker_session=sagemaker_session,
             )
         except ClientError as ce:

src/sagemaker/experiments/trial_component.py

@@ -20,6 +20,7 @@
 from sagemaker.apiutils import _base_types
 from sagemaker.experiments import _api_types
 from sagemaker.experiments._api_types import TrialComponentSearchResult
+from sagemaker.utils import format_tags
 
 
 class _TrialComponent(_base_types.Record):
@@ -191,7 +192,7 @@ def create(cls, trial_component_name, display_name=None, tags=None, sagemaker_se
         Args:
             trial_component_name (str): The name of the trial component.
             display_name (str): Display name of the trial component used by Studio (default: None).
-            tags (List[Dict[str, str]]): Tags to add to the trial component (default: None).
+            tags (Optional[Tags]): Tags to add to the trial component (default: None).
             sagemaker_session (sagemaker.session.Session): Session object which
                 manages interactions with Amazon SageMaker APIs and any other
                 AWS services needed. If not specified, one is created using the
@@ -204,7 +205,7 @@ def create(cls, trial_component_name, display_name=None, tags=None, sagemaker_se
             cls._boto_create_method,
             trial_component_name=trial_component_name,
             display_name=display_name,
-            tags=tags,
+            tags=format_tags(tags),
             sagemaker_session=sagemaker_session,
         )
 
@@ -316,7 +317,7 @@ def _load_or_create(
             display_name (str): Display name of the trial component used by Studio (default: None).
                 This is used only when the given `trial_component_name` does not
                 exist and a new trial component has to be created.
-            tags (List[Dict[str, str]]): Tags to add to the trial component (default: None).
+            tags (Optional[Tags]): Tags to add to the trial component (default: None).
                 This is used only when the given `trial_component_name` does not
                 exist and a new trial component has to be created.
             sagemaker_session (sagemaker.session.Session): Session object which
@@ -333,7 +334,7 @@ def _load_or_create(
             run_tc = _TrialComponent.create(
                 trial_component_name=trial_component_name,
                 display_name=display_name,
-                tags=tags,
+                tags=format_tags(tags),
                 sagemaker_session=sagemaker_session,
             )
         except ClientError as ce:

src/sagemaker/huggingface/model.py

@@ -29,7 +29,7 @@
 from sagemaker.predictor import Predictor
 from sagemaker.serializers import JSONSerializer
 from sagemaker.session import Session
-from sagemaker.utils import to_string
+from sagemaker.utils import to_string, format_tags
 from sagemaker.workflow import is_pipeline_variable
 from sagemaker.workflow.entities import PipelineVariable
 
@@ -255,7 +255,7 @@ def deploy(
                 https://docs.aws.amazon.com/sagemaker/latest/dg/ei.html
             endpoint_name (str): The name of the endpoint to create (default:
                 None). If not specified, a unique endpoint name will be created.
-            tags (List[dict[str, str]]): The list of tags to attach to this
+            tags (Optional[Tags]): The list of tags to attach to this
                 specific endpoint.
             kms_key (str): The ARN of the KMS key that is used to encrypt the
                 data on the storage volume attached to the instance hosting the
@@ -319,7 +319,7 @@ def deploy(
             deserializer,
             accelerator_type,
             endpoint_name,
-            tags,
+            format_tags(tags),
             kms_key,
             wait,
             data_capture_config,

src/sagemaker/jumpstart/model.py

@@ -388,7 +388,7 @@ def _create_sagemaker_model(
                 attach to an endpoint for model loading and inference, for
                 example, 'ml.eia1.medium'. If not specified, no Elastic
                 Inference accelerator will be attached to the endpoint. (Default: None).
-            tags (List[dict[str, str]]): Optional. The list of tags to add to
+            tags (Optional[Tags]): Optional. The list of tags to add to
                 the model. Example: >>> tags = [{'Key': 'tagname', 'Value':
                 'tagvalue'}] For more information about tags, see
                 https://boto3.amazonaws.com/v1/documentation
@@ -402,6 +402,8 @@ def _create_sagemaker_model(
                 any so they are ignored.
         """
 
+        tags = format_tags(tags)
+
         # if the user inputs a model artifact uri, do not use model package arn to create
         # inference endpoint.
         if self.model_package_arn and not self._model_data_is_set:

src/sagemaker/lineage/action.py

@@ -21,6 +21,7 @@
 from sagemaker.lineage import _api_types, _utils
 from sagemaker.lineage._api_types import ActionSource, ActionSummary
 from sagemaker.lineage.artifact import Artifact
+from sagemaker.utils import format_tags
 
 from sagemaker.lineage.query import (
     LineageQuery,
@@ -159,12 +160,12 @@ def set_tags(self, tags=None):
         """Add tags to the object.
 
         Args:
-            tags ([{key:value}]): list of key value pairs.
+            tags (Optional[Tags]): list of key value pairs.
 
         Returns:
             list({str:str}): a list of key value pairs
         """
-        return self._set_tags(resource_arn=self.action_arn, tags=tags)
+        return self._set_tags(resource_arn=self.action_arn, tags=format_tags(tags))
 
     @classmethod
     def create(

src/sagemaker/lineage/artifact.py

@@ -31,7 +31,7 @@
 )
 from sagemaker.lineage._utils import _disassociate, get_resource_name_from_arn
 from sagemaker.lineage.association import Association
-from sagemaker.utils import get_module
+from sagemaker.utils import get_module, format_tags
 
 LOGGER = logging.getLogger("sagemaker")
 
@@ -288,12 +288,12 @@ def set_tags(self, tags=None):
         """Add tags to the object.
 
         Args:
-            tags ([{key:value}]): list of key value pairs.
+            tags (Optional[Tags]): list of key value pairs.
 
         Returns:
             list({str:str}): a list of key value pairs
         """
-        return self._set_tags(resource_arn=self.artifact_arn, tags=tags)
+        return self._set_tags(resource_arn=self.artifact_arn, tags=format_tags(tags))
 
     @classmethod
     def create(

src/sagemaker/lineage/association.py

@@ -20,6 +20,7 @@
 from sagemaker.apiutils import _base_types
 from sagemaker.lineage import _api_types
 from sagemaker.lineage._api_types import AssociationSummary
+from sagemaker.utils import format_tags
 
 logger = logging.getLogger(__name__)
 
@@ -95,7 +96,7 @@ def set_tags(self, tags=None):
             "set_tags on Association is deprecated. Use set_tags on the source or destination\
             entity instead."
         )
-        return self._set_tags(resource_arn=self.source_arn, tags=tags)
+        return self._set_tags(resource_arn=self.source_arn, tags=format_tags(tags))
 
     @classmethod
     def create(

src/sagemaker/lineage/context.py

@@ -33,6 +33,7 @@
 from sagemaker.lineage.artifact import Artifact
 from sagemaker.lineage.action import Action
 from sagemaker.lineage.lineage_trial_component import LineageTrialComponent
+from sagemaker.utils import format_tags
 
 
 class Context(_base_types.Record):
@@ -126,7 +127,7 @@ def set_tags(self, tags=None):
         Returns:
             list({str:str}): a list of key value pairs
         """
-        return self._set_tags(resource_arn=self.context_arn, tags=tags)
+        return self._set_tags(resource_arn=self.context_arn, tags=format_tags(tags))
 
     @classmethod
     def load(cls, context_name: str, sagemaker_session=None) -> "Context":

src/sagemaker/local/entities.py

@@ -28,7 +28,7 @@
 
 from sagemaker.local.image import _SageMakerContainer
 from sagemaker.local.utils import copy_directory_structure, move_to_destination, get_docker_host
-from sagemaker.utils import DeferredError, get_config_value
+from sagemaker.utils import DeferredError, get_config_value, format_tags
 from sagemaker.local.exceptions import StepExecutionException
 
 logger = logging.getLogger(__name__)
@@ -552,7 +552,7 @@ class _LocalEndpointConfig(object):
     def __init__(self, config_name, production_variants, tags=None):
         self.name = config_name
         self.production_variants = production_variants
-        self.tags = tags
+        self.tags = format_tags(tags)
         self.creation_time = datetime.datetime.now()
 
     def describe(self):
@@ -584,7 +584,7 @@ def __init__(self, endpoint_name, endpoint_config_name, tags=None, local_session
         self.name = endpoint_name
         self.endpoint_config = local_client.describe_endpoint_config(endpoint_config_name)
         self.production_variant = self.endpoint_config["ProductionVariants"][0]
-        self.tags = tags
+        self.tags = format_tags(tags)
 
         model_name = self.production_variant["ModelName"]
         self.primary_container = local_client.describe_model(model_name)["PrimaryContainer"]