fix: jumpstart amt tracking

Author: evakravi

src/sagemaker/tuner.py

@@ -32,6 +32,7 @@
 from sagemaker.estimator import Framework
 from sagemaker.inputs import TrainingInput
 from sagemaker.job import _Job
+from sagemaker.jumpstart.utils import add_jumpstart_tags, get_jumpstart_base_name_if_jumpstart_model
 from sagemaker.parameter import (
     CategoricalParameter,
     ContinuousParameter,
@@ -318,6 +319,44 @@ def _prepare_for_tuning(self, job_name=None, include_cls_metadata=False):
         """Prepare the tuner instance for tuning (fit)."""
         self._prepare_job_name_for_tuning(job_name=job_name)
         self._prepare_static_hyperparameters_for_tuning(include_cls_metadata=include_cls_metadata)
+        self._prepare_tags_for_tuning()
+
+    def _get_model_uri(
+        self,
+        estimator,
+    ):
+        """Get model uri attribute of ``Estimator`` object.
+
+        This attribute can live in multiple places, and accessing the attribute can
+        raise a TypeError, which needs to be handled.
+        """
+        try:
+            model_data = getattr(estimator, "model_data", None)
+        except TypeError:
+            model_data = None
+        return model_data or getattr(estimator, "model_uri", None)
+
+    def _prepare_tags_for_tuning(self):
+        """Add tags to tuning job (from Estimator and JumpStart tags)."""
+
+        # Add tags from Estimator class
+        estimator = self.estimator or self.estimator_dict[sorted(self.estimator_dict.keys())[0]]
+
+        estimator_tags = getattr(estimator, "tags", []) or []
+
+        if self.tags is None and len(estimator_tags) > 0:
+            self.tags = []
+
+        for tag in estimator_tags:
+            if tag not in self.tags:
+                self.tags.append(tag)
+
+        # Add JumpStart tags
+        self.tags = add_jumpstart_tags(
+            tags=self.tags,
+            training_script_uri=getattr(estimator, "source_dir", None),
+            training_model_uri=self._get_model_uri(estimator),
+        )
 
     def _prepare_job_name_for_tuning(self, job_name=None):
         """Set current job name before starting tuning."""
@@ -330,6 +369,12 @@ def _prepare_job_name_for_tuning(self, job_name=None):
                     self.estimator or self.estimator_dict[sorted(self.estimator_dict.keys())[0]]
                 )
                 base_name = base_name_from_image(estimator.training_image_uri())
+
+                jumpstart_base_name = get_jumpstart_base_name_if_jumpstart_model(
+                    getattr(estimator, "source_dir", None),
+                    self._get_model_uri(estimator),
+                )
+                base_name = jumpstart_base_name or base_name
             self._current_job_name = name_from_base(
                 base_name, max_length=self.TUNING_JOB_NAME_MAX_LENGTH, short=True
             )

tests/unit/test_tuner.py

@@ -22,6 +22,9 @@
 from sagemaker import Predictor, TrainingInput, utils
 from sagemaker.amazon.amazon_estimator import RecordSet
 from sagemaker.estimator import Framework
+from sagemaker.fw_utils import UploadedCode
+from sagemaker.jumpstart.constants import JUMPSTART_BUCKET_NAME_SET, JUMPSTART_RESOURCE_BASE_NAME
+from sagemaker.jumpstart.enums import JumpStartTag
 from sagemaker.mxnet import MXNet
 from sagemaker.parameter import ParameterRange
 from sagemaker.tuner import (
@@ -1518,3 +1521,65 @@ def _convert_tuning_job_details(job_details, estimator_name):
     job_details_copy["TrainingJobDefinitions"] = [training_details]
 
     return job_details_copy
+
+
+@patch("time.time", return_value=510006209.073025)
+@patch("sagemaker.estimator.tar_and_upload_dir")
+@patch("sagemaker.model.Model._upload_code")
+def test_tags_prefixes_jumpstart_models(
+    patched_upload_code, patched_tar_and_upload_dir, sagemaker_session
+):
+
+    patched_tar_and_upload_dir.return_value = UploadedCode(
+        s3_prefix="s3://%s/%s" % ("bucket", "key"), script_name="script_name"
+    )
+    sagemaker_session.boto_region_name = REGION
+
+    instance_type = "ml.p2.xlarge"
+    instance_count = 1
+
+    training_data_uri = "s3://bucket/mydata"
+
+    jumpstart_source_dir = f"s3://{list(JUMPSTART_BUCKET_NAME_SET)[0]}/source_dirs/source.tar.gz"
+    jumpstart_source_dir_2 = f"s3://{list(JUMPSTART_BUCKET_NAME_SET)[1]}/source_dirs/source.tar.gz"
+
+    image_uri = "fake-image-uri"
+
+    generic_estimator = Estimator(
+        entry_point="transfer_learning.py",
+        role=ROLE,
+        region=REGION,
+        sagemaker_session=sagemaker_session,
+        instance_count=instance_count,
+        instance_type=instance_type,
+        source_dir=jumpstart_source_dir,
+        image_uri=image_uri,
+        model_uri=jumpstart_source_dir_2,
+        tags=[{"Key": "estimator-tag-key", "Value": "estimator-tag-value"}],
+    )
+
+    hp_tuner = HyperparameterTuner(
+        generic_estimator,
+        OBJECTIVE_METRIC_NAME,
+        HYPERPARAMETER_RANGES,
+        tags=[{"Key": "hp-tuner-tag-key", "Value": "hp-tuner-estimator-tag-value"}],
+    )
+
+    hp_tuner.fit({"training": training_data_uri})
+
+    assert [
+        {"Key": "hp-tuner-tag-key", "Value": "hp-tuner-estimator-tag-value"},
+        {"Key": "estimator-tag-key", "Value": "estimator-tag-value"},
+        {
+            "Key": JumpStartTag.TRAINING_MODEL_URI.value,
+            "Value": jumpstart_source_dir_2,
+        },
+        {
+            "Key": JumpStartTag.TRAINING_SCRIPT_URI.value,
+            "Value": jumpstart_source_dir,
+        },
+    ] == sagemaker_session.create_tuning_job.call_args_list[0][1]["tags"]
+
+    assert sagemaker_session.create_tuning_job.call_args_list[0][1]["job_name"].startswith(
+        JUMPSTART_RESOURCE_BASE_NAME
+    )