Change: More pythonic tags

Author: martinRenou

src/sagemaker/algorithm.py

@@ -28,6 +28,7 @@
 from sagemaker.session import Session
 from sagemaker.workflow.entities import PipelineVariable
 from sagemaker.workflow.pipeline_context import runnable_by_pipeline
+from sagemaker.utils import format_tags, Tags
 
 from sagemaker.workflow import is_pipeline_variable
 
@@ -58,7 +59,7 @@ def __init__(
         base_job_name: Optional[str] = None,
         sagemaker_session: Optional[Session] = None,
         hyperparameters: Optional[Dict[str, Union[str, PipelineVariable]]] = None,
-        tags: Optional[List[Dict[str, Union[str, PipelineVariable]]]] = None,
+        tags: Optional[Tags] = None,
         subnets: Optional[List[Union[str, PipelineVariable]]] = None,
         security_group_ids: Optional[List[Union[str, PipelineVariable]]] = None,
         model_uri: Optional[str] = None,
@@ -121,7 +122,7 @@ def __init__(
                 interactions with Amazon SageMaker APIs and any other AWS services needed. If
                 not specified, the estimator creates one using the default
                 AWS configuration chain.
-            tags (list[dict[str, str] or list[dict[str, PipelineVariable]]): List of tags for
+            tags (Union[Tags]): Tags for
                 labeling a training job. For more, see
                 https://docs.aws.amazon.com/sagemaker/latest/dg/API_Tag.html.
             subnets (list[str] or list[PipelineVariable]): List of subnet ids. If not specified
@@ -170,7 +171,7 @@ def __init__(
             output_kms_key=output_kms_key,
             base_job_name=base_job_name,
             sagemaker_session=sagemaker_session,
-            tags=tags,
+            tags=format_tags(tags),
             subnets=subnets,
             security_group_ids=security_group_ids,
             model_uri=model_uri,

src/sagemaker/automl/automl.py

@@ -28,7 +28,7 @@
 )
 from sagemaker.job import _Job
 from sagemaker.session import Session
-from sagemaker.utils import name_from_base, resolve_value_from_config
+from sagemaker.utils import name_from_base, resolve_value_from_config, format_tags, Tags
 from sagemaker.workflow.entities import PipelineVariable
 from sagemaker.workflow.pipeline_context import runnable_by_pipeline
 
@@ -127,7 +127,7 @@ def __init__(
         total_job_runtime_in_seconds: Optional[int] = None,
         job_objective: Optional[Dict[str, str]] = None,
         generate_candidate_definitions_only: Optional[bool] = False,
-        tags: Optional[List[Dict[str, str]]] = None,
+        tags: Optional[Tags] = None,
         content_type: Optional[str] = None,
         s3_data_type: Optional[str] = None,
         feature_specification_s3_uri: Optional[str] = None,
@@ -167,8 +167,7 @@ def __init__(
                 In the format of: {"MetricName": str}
             generate_candidate_definitions_only (bool): Whether to generates
                 possible candidates without training the models.
-            tags (List[dict[str, str]]): The list of tags to attach to this
-                specific endpoint.
+            tags (Optional[Tags]): Tags to attach to this specific endpoint.
             content_type (str): The content type of the data from the input source.
             s3_data_type (str): The data type for S3 data source.
                 Valid values: ManifestFile or S3Prefix.
@@ -203,7 +202,7 @@ def __init__(
         self.target_attribute_name = target_attribute_name
         self.job_objective = job_objective
         self.generate_candidate_definitions_only = generate_candidate_definitions_only
-        self.tags = tags
+        self.tags = format_tags(tags)
         self.content_type = content_type
         self.s3_data_type = s3_data_type
         self.feature_specification_s3_uri = feature_specification_s3_uri

src/sagemaker/clarify.py

@@ -33,6 +33,7 @@
 from sagemaker.session import Session
 from sagemaker.network import NetworkConfig
 from sagemaker.processing import ProcessingInput, ProcessingOutput, Processor
+from sagemaker.utils import format_tags, Tags
 
 logger = logging.getLogger(__name__)
 
@@ -1417,7 +1418,7 @@ def __init__(
         max_runtime_in_seconds: Optional[int] = None,
         sagemaker_session: Optional[Session] = None,
         env: Optional[Dict[str, str]] = None,
-        tags: Optional[List[Dict[str, str]]] = None,
+        tags: Optional[Tags] = None,
         network_config: Optional[NetworkConfig] = None,
         job_name_prefix: Optional[str] = None,
         version: Optional[str] = None,
@@ -1454,7 +1455,7 @@ def __init__(
                 using the default AWS configuration chain.
             env (dict[str, str]): Environment variables to be passed to
                 the processing jobs (default: None).
-            tags (list[dict]): List of tags to be passed to the processing job
+            tags (Optional[Tags]): Tags to be passed to the processing job
                 (default: None). For more, see
                 https://docs.aws.amazon.com/sagemaker/latest/dg/API_Tag.html.
             network_config (:class:`~sagemaker.network.NetworkConfig`):
@@ -1482,7 +1483,7 @@ def __init__(
             None,  # We set method-specific job names below.
             sagemaker_session,
             env,
-            tags,
+            format_tags(tags),
             network_config,
         )
 

src/sagemaker/estimator.py

@@ -98,6 +98,8 @@
     to_string,
     check_and_get_run_experiment_config,
     resolve_value_from_config,
+    format_tags,
+    Tags,
 )
 from sagemaker.workflow import is_pipeline_variable
 from sagemaker.workflow.entities import PipelineVariable
@@ -144,7 +146,7 @@ def __init__(
         output_kms_key: Optional[Union[str, PipelineVariable]] = None,
         base_job_name: Optional[str] = None,
         sagemaker_session: Optional[Session] = None,
-        tags: Optional[List[Dict[str, Union[str, PipelineVariable]]]] = None,
+        tags: Optional[Tags] = None,
         subnets: Optional[List[Union[str, PipelineVariable]]] = None,
         security_group_ids: Optional[List[Union[str, PipelineVariable]]] = None,
         model_uri: Optional[str] = None,
@@ -269,8 +271,8 @@ def __init__(
                 manages interactions with Amazon SageMaker APIs and any other
                 AWS services needed. If not specified, the estimator creates one
                 using the default AWS configuration chain.
-            tags (list[dict[str, str] or list[dict[str, PipelineVariable]]):
-                List of tags for labeling a training job. For more, see
+            tags (Optional[Tags]):
+                Tags for labeling a training job. For more, see
                 https://docs.aws.amazon.com/sagemaker/latest/dg/API_Tag.html.
             subnets (list[str] or list[PipelineVariable]): List of subnet ids. If not
                 specified training job will be created without VPC config.
@@ -601,6 +603,7 @@ def __init__(
         else:
             self.sagemaker_session = sagemaker_session or Session()
 
+        tags = format_tags(tags)
         self.tags = (
             add_jumpstart_uri_tags(
                 tags=tags, training_model_uri=self.model_uri, training_script_uri=self.source_dir
@@ -1347,7 +1350,7 @@ def compile_model(
         framework=None,
         framework_version=None,
         compile_max_run=15 * 60,
-        tags=None,
+        tags: Optional[Tags] = None,
         target_platform_os=None,
         target_platform_arch=None,
         target_platform_accelerator=None,
@@ -1373,7 +1376,7 @@ def compile_model(
             compile_max_run (int): Timeout in seconds for compilation (default:
                 15 * 60). After this amount of time Amazon SageMaker Neo
                 terminates the compilation job regardless of its current status.
-            tags (list[dict]): List of tags for labeling a compilation job. For
+            tags (list[dict]): Tags for labeling a compilation job. For
                 more, see
                 https://docs.aws.amazon.com/sagemaker/latest/dg/API_Tag.html.
             target_platform_os (str): Target Platform OS, for example: 'LINUX'.
@@ -1415,7 +1418,7 @@ def compile_model(
             input_shape,
             output_path,
             self.role,
-            tags,
+            format_tags(tags),
             self._compilation_job_name(),
             compile_max_run,
             framework=framework,
@@ -1527,7 +1530,7 @@ def deploy(
         model_name=None,
         kms_key=None,
         data_capture_config=None,
-        tags=None,
+        tags: Optional[Tags] = None,
         serverless_inference_config=None,
         async_inference_config=None,
         volume_size=None,
@@ -1596,8 +1599,10 @@ def deploy(
                 empty object passed through, will use pre-defined values in
                 ``ServerlessInferenceConfig`` class to deploy serverless endpoint. Deploy an
                 instance based endpoint if it's None. (default: None)
-            tags(List[dict[str, str]]): Optional. The list of tags to attach to this specific
+            tags(Optional[Tags]): Optional. Tags to attach to this specific
                 endpoint. Example:
+                >>> tags = {'tagname', 'tagvalue'}
+                Or
                 >>> tags = [{'Key': 'tagname', 'Value': 'tagvalue'}]
                 For more information about tags, see
                 https://boto3.amazonaws.com/v1/documentation\
@@ -1659,7 +1664,7 @@ def deploy(
         model.name = model_name
 
         tags = update_inference_tags_with_jumpstart_training_tags(
-            inference_tags=tags, training_tags=self.tags
+            inference_tags=format_tags(tags), training_tags=self.tags
         )
 
         return model.deploy(
@@ -2007,7 +2012,7 @@ def transformer(
         env=None,
         max_concurrent_transforms=None,
         max_payload=None,
-        tags=None,
+        tags: Optional[Tags] = None,
         role=None,
         volume_kms_key=None,
         vpc_config_override=vpc_utils.VPC_CONFIG_DEFAULT,
@@ -2041,7 +2046,7 @@ def transformer(
                 to be made to each individual transform container at one time.
             max_payload (int): Maximum size of the payload in a single HTTP
                 request to the container in MB.
-            tags (list[dict]): List of tags for labeling a transform job. If
+            tags (Optional[Tags]): Tags for labeling a transform job. If
                 none specified, then the tags used for the training job are used
                 for the transform job.
             role (str): The ``ExecutionRoleArn`` IAM Role ARN for the ``Model``,
@@ -2068,7 +2073,7 @@ def transformer(
                 model. If not specified, the estimator generates a default job name
                 based on the training image name and current timestamp.
         """
-        tags = tags or self.tags
+        tags = format_tags(tags) or self.tags
         model_name = self._get_or_create_name(model_name)
 
         if self.latest_training_job is None:
@@ -2661,7 +2666,7 @@ def __init__(
         base_job_name: Optional[str] = None,
         sagemaker_session: Optional[Session] = None,
         hyperparameters: Optional[Dict[str, Union[str, PipelineVariable]]] = None,
-        tags: Optional[List[Dict[str, Union[str, PipelineVariable]]]] = None,
+        tags: Optional[Tags] = None,
         subnets: Optional[List[Union[str, PipelineVariable]]] = None,
         security_group_ids: Optional[List[Union[str, PipelineVariable]]] = None,
         model_uri: Optional[str] = None,
@@ -2790,7 +2795,7 @@ def __init__(
                     hyperparameters. SageMaker rejects the training job request and returns an
                     validation error for detected credentials, if such user input is found.
 
-            tags (list[dict[str, str] or list[dict[str, PipelineVariable]]): List of tags for
+            tags (Optional[Tags]): Tags for
                 labeling a training job. For more, see
                 https://docs.aws.amazon.com/sagemaker/latest/dg/API_Tag.html.
             subnets (list[str] or list[PipelineVariable]): List of subnet ids.
@@ -3071,7 +3076,7 @@ def __init__(
             output_kms_key,
             base_job_name,
             sagemaker_session,
-            tags,
+            format_tags(tags),
             subnets,
             security_group_ids,
             model_uri=model_uri,
@@ -3702,7 +3707,7 @@ def transformer(
         env=None,
         max_concurrent_transforms=None,
         max_payload=None,
-        tags=None,
+        tags: Optional[Tags] = None,
         role=None,
         model_server_workers=None,
         volume_kms_key=None,
@@ -3738,7 +3743,7 @@ def transformer(
                 to be made to each individual transform container at one time.
             max_payload (int): Maximum size of the payload in a single HTTP
                 request to the container in MB.
-            tags (list[dict]): List of tags for labeling a transform job. If
+            tags (Optional[Tags]): Tags for labeling a transform job. If
                 none specified, then the tags used for the training job are used
                 for the transform job.
             role (str): The ``ExecutionRoleArn`` IAM Role ARN for the ``Model``,
@@ -3777,7 +3782,7 @@ def transformer(
                 SageMaker Batch Transform job.
         """
         role = role or self.role
-        tags = tags or self.tags
+        tags = format_tags(tags) or self.tags
         model_name = self._get_or_create_name(model_name)
 
         if self.latest_training_job is not None:

src/sagemaker/experiments/run.py

@@ -44,6 +44,9 @@
 from sagemaker.utils import (
     get_module,
     unique_name_from_base,
+    format_tags,
+    Tags,
+    TagsDict,
 )
 
 from sagemaker.experiments._utils import (
@@ -97,7 +100,7 @@ def __init__(
         run_name: Optional[str] = None,
         experiment_display_name: Optional[str] = None,
         run_display_name: Optional[str] = None,
-        tags: Optional[List[Dict[str, str]]] = None,
+        tags: Optional[Tags] = None,
         sagemaker_session: Optional["Session"] = None,
         artifact_bucket: Optional[str] = None,
         artifact_prefix: Optional[str] = None,
@@ -152,7 +155,7 @@ def __init__(
             run_display_name (str): The display name of the run used in UI (default: None).
                 This display name is used in a create run call. If a run with the
                 specified name already exists, this display name won't take effect.
-            tags (List[Dict[str, str]]): A list of tags to be used for all create calls,
+            tags (Optional[Tags]): Tags to be used for all create calls,
                 e.g. to create an experiment, a run group, etc. (default: None).
             sagemaker_session (sagemaker.session.Session): Session object which
                 manages interactions with Amazon SageMaker APIs and any other
@@ -172,6 +175,8 @@ def __init__(
         # avoid confusion due to mis-match in casing between run name and TC name
         self.run_name = self.run_name.lower()
 
+        tags = format_tags(tags)
+
         trial_component_name = Run._generate_trial_component_name(
             run_name=self.run_name, experiment_name=self.experiment_name
         )
@@ -676,11 +681,11 @@ def _extract_run_name_from_tc_name(trial_component_name: str, experiment_name: s
         )
 
     @staticmethod
-    def _append_run_tc_label_to_tags(tags: Optional[List[Dict[str, str]]] = None) -> list:
+    def _append_run_tc_label_to_tags(tags: Optional[List[TagsDict]] = None) -> list:
         """Append the run trial component label to tags used to create a trial component.
 
         Args:
-            tags (List[Dict[str, str]]): The tags supplied by users to initialize a Run object.
+            tags (List[TagsDict]): The tags supplied by users to initialize a Run object.
 
         Returns:
             list: The updated tags with the appended run trial component label.

src/sagemaker/feature_store/feature_group.py

@@ -28,7 +28,7 @@
 import tempfile
 from concurrent.futures import as_completed
 from concurrent.futures import ThreadPoolExecutor
-from typing import Sequence, List, Dict, Any, Union
+from typing import Optional, Sequence, List, Dict, Any, Union
 from urllib.parse import urlparse
 
 from multiprocessing.pool import AsyncResult
@@ -65,7 +65,7 @@
     OnlineStoreConfigUpdate,
     OnlineStoreStorageTypeEnum,
 )
-from sagemaker.utils import resolve_value_from_config
+from sagemaker.utils import resolve_value_from_config, format_tags, Tags
 
 logger = logging.getLogger(__name__)
 
@@ -538,7 +538,7 @@ def create(
         disable_glue_table_creation: bool = False,
         data_catalog_config: DataCatalogConfig = None,
         description: str = None,
-        tags: List[Dict[str, str]] = None,
+        tags: Optional[Tags] = None,
         table_format: TableFormatEnum = None,
         online_store_storage_type: OnlineStoreStorageTypeEnum = None,
     ) -> Dict[str, Any]:
@@ -566,7 +566,7 @@ def create(
             data_catalog_config (DataCatalogConfig): configuration for
                 Metadata store (default: None).
             description (str): description of the FeatureGroup (default: None).
-            tags (List[Dict[str, str]]): list of tags for labeling a FeatureGroup (default: None).
+            tags (Optional[Tags]): Tags for labeling a FeatureGroup (default: None).
             table_format (TableFormatEnum): format of the offline store table (default: None).
             online_store_storage_type (OnlineStoreStorageTypeEnum): storage type for the
                 online store (default: None).
@@ -602,7 +602,7 @@ def create(
             ],
             role_arn=role_arn,
             description=description,
-            tags=tags,
+            tags=format_tags(tags),
         )
 
         # online store configuration

src/sagemaker/feature_store/feature_processor/_event_bridge_rule_helper.py

@@ -32,6 +32,7 @@
 from sagemaker.feature_store.feature_processor._enums import (
     FeatureProcessorPipelineExecutionStatus,
 )
+from sagemaker.utils import TagsDict
 
 logger = logging.getLogger("sagemaker")
 
@@ -175,7 +176,7 @@ def disable_rule(self, rule_name: str) -> None:
         self.event_bridge_rule_client.disable_rule(Name=rule_name)
         logger.info("Disabled EventBridge Rule for pipeline %s.", rule_name)
 
-    def add_tags(self, rule_arn: str, tags: List[Dict[str, str]]) -> None:
+    def add_tags(self, rule_arn: str, tags: List[TagsDict]) -> None:
         """Adds tags to the EventBridge Rule.
 
         Args: