fix: make telemetry logger persist certain information (aws#1500)

* refactor telemetry logger

* refactor

* refactor

* pylint + UT

* add tag

* add remove tags

* handle tags again

* pylint

---------

Co-authored-by: Jacky Lee <[email protected]>

Author: grenmester

src/sagemaker/model.py

@@ -74,6 +74,7 @@
     Tags,
     _resolve_routing_config,
     _validate_new_tags,
+    remove_tag_with_key,
 )
 from sagemaker.async_inference import AsyncInferenceConfig
 from sagemaker.predictor_async import AsyncPredictor
@@ -426,6 +427,14 @@ def add_tags(self, tags: Tags) -> None:
         """
         self._tags = _validate_new_tags(tags, self._tags)
 
+    def remove_tag_with_key(self, key: str) -> None:
+        """Remove a tag with the given key from the list of tags.
+
+        Args:
+            key (str): The key of the tag to remove.
+        """
+        self._tags = remove_tag_with_key(key, self._tags)
+
     @classmethod
     def attach(
         cls,

src/sagemaker/serve/builder/jumpstart_builder.py

@@ -116,7 +116,9 @@ def __init__(self):
         self.model_metadata = None
         self.role_arn = None
         self.is_fine_tuned = None
-        self.is_gated = None
+        self.is_compiled = False
+        self.is_quantized = False
+        self.speculative_decoding_draft_model_source = None
 
     @abstractmethod
     def _prepare_for_mode(self):
@@ -503,6 +505,18 @@ def set_deployment_config(self, config_name: str, instance_type: str) -> None:
 
         self.pysdk_model.set_deployment_config(config_name, instance_type)
 
+        self.instance_type = instance_type
+
+        # JS-benchmarked models only include SageMaker-provided SD models
+        if self.pysdk_model.additional_model_data_sources:
+            self.speculative_decoding_draft_model_source = "sagemaker"
+            self.pysdk_model.add_tags(
+                {"Key": Tag.SPECULATIVE_DRAFT_MODEL_PROVIDER, "Value": "sagemaker"},
+            )
+            self.pysdk_model.remove_tag_with_key(Tag.OPTIMIZATION_JOB_NAME)
+            self.pysdk_model.remove_tag_with_key(Tag.FINE_TUNING_MODEL_PATH)
+            self.pysdk_model.remove_tag_with_key(Tag.FINE_TUNING_JOB_NAME)
+
     def get_deployment_config(self) -> Optional[Dict[str, Any]]:
         """Gets the deployment config to apply to the model.
 
@@ -775,10 +789,8 @@ def _is_gated_model(self, model=None) -> bool:
             s3_uri = s3_uri.get("S3DataSource").get("S3Uri")
 
         if s3_uri is None:
-            self.is_gated = False
-        else:
-            self.is_gated = "private" in s3_uri
-        return self.is_gated
+            return False
+        return "private" in s3_uri
 
     def _set_additional_model_source(
         self,

src/sagemaker/serve/builder/model_builder.py

@@ -23,6 +23,7 @@
 
 from pathlib import Path
 
+from sagemaker.enums import Tag
 from sagemaker.s3 import S3Downloader
 
 from sagemaker import Session
@@ -69,6 +70,7 @@
 from sagemaker.serve.utils.lineage_utils import _maintain_lineage_tracking_for_mlflow_model
 from sagemaker.serve.utils.optimize_utils import (
     _generate_optimized_model,
+    _extract_speculative_draft_model_provider,
 )
 from sagemaker.serve.utils.predictors import _get_local_mode_predictor
 from sagemaker.serve.utils.hardware_detector import (
@@ -647,11 +649,6 @@ def _handle_mlflow_input(self):
         mlflow_model_path = self.model_metadata.get(MLFLOW_MODEL_PATH)
         artifact_path = self._get_artifact_path(mlflow_model_path)
         if not self._mlflow_metadata_exists(artifact_path):
-            logger.info(
-                "MLflow model metadata not detected in %s. ModelBuilder is not "
-                "handling MLflow model input",
-                mlflow_model_path,
-            )
             return
 
         self._initialize_for_mlflow(artifact_path)
@@ -1144,6 +1141,12 @@ def _model_builder_optimize_wrapper(
         Returns:
             Model: A deployable ``Model`` object.
         """
+        self.is_compiled = compilation_config is not None
+        self.is_quantized = quantization_config is not None
+        self.speculative_decoding_draft_model_source = _extract_speculative_draft_model_provider(
+            speculative_decoding_config
+        )
+
         if quantization_config and compilation_config:
             raise ValueError("Quantization config and compilation config are mutually exclusive.")
 
@@ -1180,4 +1183,8 @@ def _model_builder_optimize_wrapper(
             job_status = self.sagemaker_session.wait_for_optimization_job(job_name)
             return _generate_optimized_model(self.pysdk_model, job_status)
 
+        self.pysdk_model.remove_tag_with_key(Tag.OPTIMIZATION_JOB_NAME)
+        if not speculative_decoding_config:
+            self.pysdk_model.remove_tag_with_key(Tag.SPECULATIVE_DRAFT_MODEL_PROVIDER)
+
         return self.pysdk_model

src/sagemaker/serve/utils/telemetry_logger.py

@@ -94,15 +94,38 @@ def wrapper(self, *args, **kwargs):
             logger.info(TELEMETRY_OPT_OUT_MESSAGING)
             response = None
             caught_ex = None
-
+            status = "1"
+            failure_reason = None
+            failure_type = None
             extra = f"{func_name}"
 
+            start_timer = perf_counter()
+            try:
+                response = func(self, *args, **kwargs)
+            except (
+                ModelBuilderException,
+                exceptions.CapacityError,
+                exceptions.UnexpectedStatusException,
+                exceptions.AsyncInferenceError,
+            ) as e:
+                status = "0"
+                caught_ex = e
+                failure_reason = str(e)
+                failure_type = e.__class__.__name__
+            except Exception as e:  # pylint: disable=W0703
+                raise e
+
+            stop_timer = perf_counter()
+            elapsed = stop_timer - start_timer
+
             if self.model_server:
                 extra += f"&x-modelServer={MODEL_SERVER_TO_CODE[str(self.model_server)]}"
 
             if self.image_uri:
                 image_uri_tail = self.image_uri.split("/")[1]
-                image_uri_option = _get_image_uri_option(self.image_uri, self._is_custom_image_uri)
+                image_uri_option = _get_image_uri_option(
+                    self.image_uri, getattr(self, "_is_custom_image_uri", False)
+                )
 
             if self.image_uri:
                 extra += f"&x-imageTag={image_uri_tail}"
@@ -128,63 +151,36 @@ def wrapper(self, *args, **kwargs):
 
             if getattr(self, "is_fine_tuned", False):
                 extra += "&x-fineTuned=1"
-            if getattr(self, "is_gated", False):
-                extra += "&x-gated=1"
 
-            if kwargs.get("compilation_config"):
+            if getattr(self, "is_compiled", False):
                 extra += "&x-compiled=1"
-            if kwargs.get("quantization_config"):
+            if getattr(self, "is_quantized", False):
                 extra += "&x-quantized=1"
-            if kwargs.get("speculative_decoding_config"):
-                model_provider = kwargs["speculative_decoding_config"]["ModelProvider"]
+            if getattr(self, "speculative_decoding_draft_model_source", False):
                 model_provider_enum = (
                     SpeculativeDecodingDraftModelSource.SAGEMAKER
-                    if model_provider.lower() == "sagemaker"
+                    if self.speculative_decoding_draft_model_source == "sagemaker"
                     else SpeculativeDecodingDraftModelSource.CUSTOM
                 )
                 model_provider_value = SD_DRAFT_MODEL_SOURCE_TO_CODE[str(model_provider_enum)]
                 extra += f"&x-sdDraftModelSource={model_provider_value}"
 
-            start_timer = perf_counter()
-            try:
-                response = func(self, *args, **kwargs)
-                stop_timer = perf_counter()
-                elapsed = stop_timer - start_timer
-                extra += f"&x-latency={round(elapsed, 2)}"
-                if not self.serve_settings.telemetry_opt_out:
-                    _send_telemetry(
-                        "1",
-                        MODE_TO_CODE[str(self.mode)],
-                        self.sagemaker_session,
-                        None,
-                        None,
-                        extra,
-                    )
-            except (
-                ModelBuilderException,
-                exceptions.CapacityError,
-                exceptions.UnexpectedStatusException,
-                exceptions.AsyncInferenceError,
-            ) as e:
-                stop_timer = perf_counter()
-                elapsed = stop_timer - start_timer
-                extra += f"&x-latency={round(elapsed, 2)}"
-                if not self.serve_settings.telemetry_opt_out:
-                    _send_telemetry(
-                        "0",
-                        MODE_TO_CODE[str(self.mode)],
-                        self.sagemaker_session,
-                        str(e),
-                        e.__class__.__name__,
-                        extra,
-                    )
-                caught_ex = e
-            except Exception as e:  # pylint: disable=W0703
-                caught_ex = e
-            finally:
-                if caught_ex:
-                    raise caught_ex
-                return response  # pylint: disable=W0150
+            extra += f"&x-latency={round(elapsed, 2)}"
+
+            if not self.serve_settings.telemetry_opt_out:
+                _send_telemetry(
+                    status,
+                    MODE_TO_CODE[str(self.mode)],
+                    self.sagemaker_session,
+                    failure_reason,
+                    failure_type,
+                    extra,
+                )
+
+            if caught_ex:
+                raise caught_ex
+
+            return response
 
         return wrapper
 

src/sagemaker/utils.py

@@ -1873,3 +1873,30 @@ def _validate_new_tags(new_tags: Optional[Tags], curr_tags: Optional[Tags]) -> O
                 curr_tags.append(new_tag)
 
     return curr_tags
+
+
+def remove_tag_with_key(key: str, tags: Optional[Tags]) -> Optional[Tags]:
+    """Remove a tag with the given key from the list of tags.
+
+    Args:
+        key (str): The key of the tag to remove.
+        tags (Optional[Tags]): The current list of tags.
+
+    Returns:
+        Optional[Tags]: The updated list of tags with the tag removed.
+    """
+    if tags is None:
+        return tags
+    if isinstance(tags, dict):
+        tags = [tags]
+
+    updated_tags = []
+    for tag in tags:
+        if tag["Key"] != key:
+            updated_tags.append(tag)
+
+    if not updated_tags:
+        return None
+    if len(updated_tags) == 1:
+        return updated_tags[0]
+    return updated_tags

tests/unit/sagemaker/serve/utils/test_telemetry_logger.py

@@ -314,17 +314,15 @@ def test_capture_telemetry_decorator_optimize_with_custom_configs(self, mock_sen
         mock_model_builder.model_server = ModelServer.TORCHSERVE
         mock_model_builder.sagemaker_session.endpoint_arn = None
         mock_model_builder.is_fine_tuned = True
-        mock_model_builder.is_gated = True
+        mock_model_builder.is_compiled = True
+        mock_model_builder.is_quantized = True
+        mock_model_builder.speculative_decoding_draft_model_source = "sagemaker"
 
         mock_speculative_decoding_config = MagicMock()
         mock_config = {"ModelProvider": "sagemaker"}
         mock_speculative_decoding_config.__getitem__.side_effect = mock_config.__getitem__
 
-        mock_model_builder.mock_optimize(
-            quantization_config=Mock(),
-            compilation_config=Mock(),
-            speculative_decoding_config=mock_speculative_decoding_config,
-        )
+        mock_model_builder.mock_optimize()
 
         args = mock_send_telemetry.call_args.args
         latency = str(args[5]).split("latency=")[1]
@@ -333,7 +331,6 @@ def test_capture_telemetry_decorator_optimize_with_custom_configs(self, mock_sen
             "&x-modelServer=1"
             f"&x-sdkVersion={SDK_VERSION}"
             f"&x-fineTuned=1"
-            f"&x-gated=1"
             f"&x-compiled=1"
             f"&x-quantized=1"
             f"&x-sdDraftModelSource=1"

tests/unit/test_utils.py

@@ -57,6 +57,7 @@
     _resolve_routing_config,
     tag_exists,
     _validate_new_tags,
+    remove_tag_with_key,
 )
 from tests.unit.sagemaker.workflow.helpers import CustomStep
 from sagemaker.workflow.parameters import ParameterString, ParameterInteger
@@ -2124,3 +2125,24 @@ def test_new_add_tags(self):
         new_tag = {"Key": "project-2", "Value": "my-project-2"}
 
         self.assertEqual(_validate_new_tags(new_tag, None), new_tag)
+
+    def test_remove_existing_tag(self):
+        original_tags = [
+            {"Key": "Tag1", "Value": "Value1"},
+            {"Key": "Tag2", "Value": "Value2"},
+            {"Key": "Tag3", "Value": "Value3"},
+        ]
+        expected_output = [{"Key": "Tag1", "Value": "Value1"}, {"Key": "Tag3", "Value": "Value3"}]
+        self.assertEqual(remove_tag_with_key("Tag2", original_tags), expected_output)
+
+    def test_remove_non_existent_tag(self):
+        original_tags = [
+            {"Key": "Tag1", "Value": "Value1"},
+            {"Key": "Tag2", "Value": "Value2"},
+            {"Key": "Tag3", "Value": "Value3"},
+        ]
+        self.assertEqual(remove_tag_with_key("NonExistentTag", original_tags), original_tags)
+
+    def test_remove_only_tag(self):
+        original_tags = [{"Key": "Tag1", "Value": "Value1"}]
+        self.assertIsNone(remove_tag_with_key("Tag1", original_tags))