change: add project tags to creates (aws#534)

Author: metrizable

src/sagemaker/_studio.py

@@ -23,7 +23,7 @@
 logger = logging.getLogger(__name__)
 
 
-def _append_project_tags(working_dir=None, tags=None):
+def _append_project_tags(tags=None, working_dir=None):
     """Appends the project tag to the list of tags, if it exists.
 
     Args:

src/sagemaker/session.py

@@ -31,6 +31,7 @@
 import sagemaker.logs
 from sagemaker import vpc_utils
 
+from sagemaker._studio import _append_project_tags
 from sagemaker.deprecations import deprecated_class
 from sagemaker.inputs import ShuffleConfig, TrainingInput
 from sagemaker.user_agent import prepend_user_agent
@@ -534,6 +535,7 @@ def train(  # noqa: C901
         Returns:
             str: ARN of the training job, if it is created.
         """
+        tags = _append_project_tags(tags)
         train_request = self._get_train_request(
             input_mode=input_mode,
             input_config=input_config,
@@ -779,6 +781,7 @@ def process(
                 three optional keys, 'ExperimentName', 'TrialName', and 'TrialComponentDisplayName'.
                 (default: ``None``)
         """
+        tags = _append_project_tags(tags)
         process_request = self._get_process_request(
             inputs=inputs,
             output_config=output_config,
@@ -1019,6 +1022,7 @@ def create_monitoring_schedule(
                 "NetworkConfig"
             ] = network_config
 
+        tags = _append_project_tags(tags)
         if tags is not None:
             monitoring_schedule_request["Tags"] = tags
 
@@ -1527,6 +1531,8 @@ def auto_ml(
             auto_ml_job_request["AutoMLJobObjective"] = job_objective
         if problem_type is not None:
             auto_ml_job_request["ProblemType"] = problem_type
+
+        tags = _append_project_tags(tags)
         if tags is not None:
             auto_ml_job_request["Tags"] = tags
 
@@ -1719,6 +1725,7 @@ def compile_model(
             "CompilationJobName": job_name,
         }
 
+        tags = _append_project_tags(tags)
         if tags is not None:
             compilation_job_request["Tags"] = tags
 
@@ -1868,6 +1875,7 @@ def tune(  # noqa: C901
         if warm_start_config is not None:
             tune_request["WarmStartConfig"] = warm_start_config
 
+        tags = _append_project_tags(tags)
         if tags is not None:
             tune_request["Tags"] = tags
 
@@ -1925,6 +1933,7 @@ def create_tuning_job(
         if warm_start_config is not None:
             tune_request["WarmStartConfig"] = warm_start_config
 
+        tags = _append_project_tags(tags)
         if tags is not None:
             tune_request["Tags"] = tags
 
@@ -2315,6 +2324,7 @@ def transform(
                 job. Dictionary contains two optional keys,
                 'InvocationsTimeoutInSeconds', and 'InvocationsMaxRetries'.
         """
+        tags = _append_project_tags(tags)
         transform_request = self._get_transform_request(
             job_name=job_name,
             model_name=model_name,
@@ -2430,6 +2440,7 @@ def create_model(
         Returns:
             str: Name of the Amazon SageMaker ``Model`` created.
         """
+        tags = _append_project_tags(tags)
         create_model_request = self._create_model_request(
             name=name,
             role=role,
@@ -2754,6 +2765,7 @@ def create_endpoint_config(
             ],
         }
 
+        tags = _append_project_tags(tags)
         if tags is not None:
             request["Tags"] = tags
 
@@ -2823,6 +2835,7 @@ def create_endpoint_config_from_existing(
         request_tags = new_tags or self.list_tags(
             existing_endpoint_config_desc["EndpointConfigArn"]
         )
+        request_tags = _append_project_tags(request_tags)
         if request_tags:
             request["Tags"] = request_tags
 
@@ -2857,6 +2870,7 @@ def create_endpoint(self, endpoint_name, config_name, tags=None, wait=True):
         LOGGER.info("Creating endpoint with name %s", endpoint_name)
 
         tags = tags or []
+        tags = _append_project_tags(tags)
 
         self.sagemaker_client.create_endpoint(
             EndpointName=endpoint_name, EndpointConfigName=config_name, Tags=tags
@@ -3336,6 +3350,7 @@ def endpoint_from_production_variants(
             lambda: self.sagemaker_client.describe_endpoint_config(EndpointConfigName=name)
         ):
             config_options = {"EndpointConfigName": name, "ProductionVariants": production_variants}
+            tags = _append_project_tags(tags)
             if tags:
                 config_options["Tags"] = tags
             if kms_key:
@@ -3728,6 +3743,7 @@ def create_feature_group(
         Returns:
             Response dict from service.
         """
+        tags = _append_project_tags(tags)
         kwargs = dict(
             FeatureGroupName=feature_group_name,
             RecordIdentifierFeatureName=record_identifier_name,

tests/unit/sagemaker/test_studio.py

@@ -82,13 +82,13 @@ def test_append_project_tags(tmpdir):
     config.write('{"sagemakerProjectId": "proj-1234", "sagemakerProjectName": "proj-name"}')
     working_dir = tmpdir.mkdir("sub")
 
-    tags = _append_project_tags(working_dir, None)
+    tags = _append_project_tags(None, working_dir)
     assert tags == [
         {"Key": "sagemaker:project-id", "Value": "proj-1234"},
         {"Key": "sagemaker:project-name", "Value": "proj-name"},
     ]
 
-    tags = _append_project_tags(working_dir, [{"Key": "a", "Value": "b"}])
+    tags = _append_project_tags([{"Key": "a", "Value": "b"}], working_dir)
     assert tags == [
         {"Key": "a", "Value": "b"},
         {"Key": "sagemaker:project-id", "Value": "proj-1234"},