feat: add quicksilver telemetry (aws#1482)

* feat: add quicksilver telemetry fields

* pylint

* add UTs

* pylint

* Refactor

* add gated and fine-tuned to telemetry

* fix: typo

* fix: jumpstart var

* refactor model_hub

* pylint

* update TEI/TGI to remove jumpstart field

* reorder telemetry schema

* refactor

---------

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

Author: grenmester

src/sagemaker/enums.py

@@ -46,6 +46,6 @@ class Tag(str, Enum):
     """Enum class for tag keys to apply to models."""
 
     OPTIMIZATION_JOB_NAME = "sagemaker-sdk:optimization-job-name"
-    SPECULATIVE_DRAFT_MODL_PROVIDER = "sagemaker-sdk:speculative-draft-model-provider"
+    SPECULATIVE_DRAFT_MODEL_PROVIDER = "sagemaker-sdk:speculative-draft-model-provider"
     FINE_TUNING_MODEL_PATH = "sagemaker-sdk:fine-tuning-model-path"
     FINE_TUNING_JOB_NAME = "sagemaker-sdk:fine-tuning-job-name"

src/sagemaker/serve/builder/jumpstart_builder.py

@@ -108,8 +108,10 @@ def __init__(self):
         self.schema_builder = None
         self.nb_instance_type = None
         self.ram_usage_model_load = None
-        self.jumpstart = None
+        self.model_hub = None
         self.model_metadata = None
+        self.is_fine_tuned = None
+        self.is_gated = None
 
     @abstractmethod
     def _prepare_for_mode(self):
@@ -580,14 +582,14 @@ def _build_for_jumpstart(self):
 
         # we do not pickle for jumpstart. set to none
         self.secret_key = None
-        self.jumpstart = True
 
         pysdk_model = self._create_pre_trained_js_model()
         image_uri = pysdk_model.image_uri
 
         logger.info("JumpStart ID %s is packaged with Image URI: %s", self.model, image_uri)
 
         if self._is_fine_tuned_model():
+            self.is_fine_tuned = True
             pysdk_model = self._update_model_data_for_fine_tuned_model(pysdk_model)
 
         if self._is_gated_model(pysdk_model) and self.mode != Mode.SAGEMAKER_ENDPOINT:
@@ -754,8 +756,10 @@ def _is_gated_model(self, model=None) -> bool:
             s3_uri = s3_uri.get("S3DataSource").get("S3Uri")
 
         if s3_uri is None:
-            return False
-        return "private" in s3_uri
+            self.is_gated = False
+        else:
+            self.is_gated = "private" in s3_uri
+        return self.is_gated
 
     def _set_additional_model_source(
         self,
@@ -792,7 +796,7 @@ def _set_additional_model_source(
                         )
 
                 self.pysdk_model.add_tags(
-                    {"key": Tag.SPECULATIVE_DRAFT_MODL_PROVIDER, "value": "sagemaker"},
+                    {"key": Tag.SPECULATIVE_DRAFT_MODEL_PROVIDER, "value": "sagemaker"},
                 )
             else:
                 s3_uri = speculative_decoding_config.get("ModelSource")
@@ -811,7 +815,7 @@ def _set_additional_model_source(
 
                 self.pysdk_model.additional_model_data_sources = [additional_model_data_source]
                 self.pysdk_model.add_tags(
-                    {"key": Tag.SPECULATIVE_DRAFT_MODL_PROVIDER, "value": "customer"},
+                    {"key": Tag.SPECULATIVE_DRAFT_MODEL_PROVIDER, "value": "customer"},
                 )
 
     def _find_compatible_deployment_config(

src/sagemaker/serve/builder/model_builder.py

@@ -78,7 +78,7 @@
 from sagemaker.serve.model_server.torchserve.prepare import prepare_for_torchserve
 from sagemaker.serve.model_server.triton.triton_builder import Triton
 from sagemaker.serve.utils.telemetry_logger import _capture_telemetry
-from sagemaker.serve.utils.types import ModelServer
+from sagemaker.serve.utils.types import ModelServer, ModelHub
 from sagemaker.serve.validations.check_image_uri import is_1p_image_uri
 from sagemaker.serve.save_retrive.version_1_0_0.save.save_handler import SaveHandler
 from sagemaker.serve.save_retrive.version_1_0_0.metadata.metadata import get_metadata
@@ -400,7 +400,7 @@ def _prepare_for_mode(self):
                 self.serve_settings.s3_model_data_url,
                 self.sagemaker_session,
                 self.image_uri,
-                self.jumpstart if hasattr(self, "jumpstart") else False,
+                getattr(self, "model_hub", None) == ModelHub.JUMPSTART,
             )
             self.env_vars.update(env_vars_sagemaker)
             return self.s3_upload_path, env_vars_sagemaker
@@ -754,10 +754,14 @@ def build(  # pylint: disable=R0911
 
         if isinstance(self.model, str):
             model_task = None
-            if self.model_metadata:
-                model_task = self.model_metadata.get("HF_TASK")
             if self._is_jumpstart_model_id():
+                self.model_hub = ModelHub.JUMPSTART
                 return self._build_for_jumpstart()
+            self.model_hub = ModelHub.HUGGINGFACE
+
+            if self.model_metadata:
+                model_task = self.model_metadata.get("HF_TASK")
+
             if self._is_djl():
                 return self._build_for_djl()
             else:

src/sagemaker/serve/builder/tei_builder.py

@@ -63,7 +63,6 @@ def __init__(self):
         self.nb_instance_type = None
         self.ram_usage_model_load = None
         self.secret_key = None
-        self.jumpstart = None
         self.role_arn = None
 
     @abstractmethod

src/sagemaker/serve/builder/tgi_builder.py

@@ -90,7 +90,6 @@ def __init__(self):
         self.nb_instance_type = None
         self.ram_usage_model_load = None
         self.secret_key = None
-        self.jumpstart = None
         self.role_arn = None
 
     @abstractmethod

src/sagemaker/serve/utils/telemetry_logger.py

@@ -29,7 +29,12 @@
     MLFLOW_REGISTRY_PATH,
 )
 from sagemaker.serve.utils.lineage_utils import _get_mlflow_model_path_type
-from sagemaker.serve.utils.types import ModelServer, ImageUriOption
+from sagemaker.serve.utils.types import (
+    ModelServer,
+    ImageUriOption,
+    ModelHub,
+    SpeculativeDecodingDraftModelSource,
+)
 from sagemaker.serve.validations.check_image_uri import is_1p_image_uri
 from sagemaker.user_agent import SDK_VERSION
 
@@ -69,6 +74,16 @@
     MLFLOW_REGISTRY_PATH: 5,
 }
 
+MODEL_HUB_TO_CODE = {
+    str(ModelHub.JUMPSTART): 1,
+    str(ModelHub.HUGGINGFACE): 2,
+}
+
+SD_DRAFT_MODEL_SOURCE_TO_CODE = {
+    str(SpeculativeDecodingDraftModelSource.SAGEMAKER): 1,
+    str(SpeculativeDecodingDraftModelSource.CUSTOM): 2,
+}
+
 
 def _capture_telemetry(func_name: str):
     """Placeholder docstring"""
@@ -108,6 +123,28 @@ def wrapper(self, *args, **kwargs):
                 mlflow_model_path_type = _get_mlflow_model_path_type(mlflow_model_path)
                 extra += f"&x-mlflowModelPathType={MLFLOW_MODEL_PATH_CODE[mlflow_model_path_type]}"
 
+            if getattr(self, "model_hub", False):
+                extra += f"&x-modelHub={MODEL_HUB_TO_CODE[str(self.model_hub)]}"
+
+            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"):
+                extra += "&x-compiled=1"
+            if kwargs.get("quantization_config"):
+                extra += "&x-quantized=1"
+            if kwargs.get("speculative_decoding_config"):
+                model_provider = kwargs["speculative_decoding_config"]["ModelProvider"]
+                model_provider_enum = (
+                    SpeculativeDecodingDraftModelSource.SAGEMAKER
+                    if model_provider.lower() == "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)

src/sagemaker/serve/utils/types.py

@@ -57,3 +57,25 @@ def __str__(self) -> str:
     CUSTOM_IMAGE = 1
     CUSTOM_1P_IMAGE = 2
     DEFAULT_IMAGE = 3
+
+
+class ModelHub(Enum):
+    """Enum type for model hub source"""
+
+    def __str__(self) -> str:
+        """Convert enum to string"""
+        return str(self.name)
+
+    JUMPSTART = 1
+    HUGGINGFACE = 2
+
+
+class SpeculativeDecodingDraftModelSource(Enum):
+    """Enum type for speculative decoding draft model source"""
+
+    def __str__(self) -> str:
+        """Convert enum to string"""
+        return str(self.name)
+
+    SAGEMAKER = 1
+    CUSTOM = 2

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

@@ -12,7 +12,7 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 import unittest
-from unittest.mock import Mock, patch
+from unittest.mock import Mock, patch, MagicMock
 from sagemaker.serve import Mode, ModelServer
 from sagemaker.serve.model_format.mlflow.constants import MLFLOW_MODEL_PATH
 from sagemaker.serve.utils.telemetry_logger import (
@@ -25,7 +25,8 @@
 from sagemaker.user_agent import SDK_VERSION
 
 MOCK_SESSION = Mock()
-MOCK_FUNC_NAME = "Mock.deploy"
+MOCK_DEPLOY_FUNC_NAME = "Mock.deploy"
+MOCK_OPTIMIZE_FUNC_NAME = "Mock.optimize"
 MOCK_DJL_CONTAINER = (
     "763104351884.dkr.ecr.us-west-2.amazonaws.com/" "djl-inference:0.25.0-deepspeed0.11.0-cu118"
 )
@@ -47,11 +48,15 @@ def __init__(self):
         self.serve_settings = Mock()
         self.sagemaker_session = MOCK_SESSION
 
-    @_capture_telemetry(MOCK_FUNC_NAME)
+    @_capture_telemetry(MOCK_DEPLOY_FUNC_NAME)
     def mock_deploy(self, mock_exception_func=None):
         if mock_exception_func:
             mock_exception_func()
 
+    @_capture_telemetry(MOCK_OPTIMIZE_FUNC_NAME)
+    def mock_optimize(self, *args, **kwargs):
+        pass
+
 
 class TestTelemetryLogger(unittest.TestCase):
     @patch("sagemaker.serve.utils.telemetry_logger._requests_helper")
@@ -88,7 +93,7 @@ def test_capture_telemetry_decorator_djl_success(self, mock_send_telemetry):
         args = mock_send_telemetry.call_args.args
         latency = str(args[5]).split("latency=")[1]
         expected_extra_str = (
-            f"{MOCK_FUNC_NAME}"
+            f"{MOCK_DEPLOY_FUNC_NAME}"
             "&x-modelServer=4"
             "&x-imageTag=djl-inference:0.25.0-deepspeed0.11.0-cu118"
             f"&x-sdkVersion={SDK_VERSION}"
@@ -118,7 +123,7 @@ def test_capture_telemetry_decorator_djl_success_with_custom_image(self, mock_se
         args = mock_send_telemetry.call_args.args
         latency = str(args[5]).split("latency=")[1]
         expected_extra_str = (
-            f"{MOCK_FUNC_NAME}"
+            f"{MOCK_DEPLOY_FUNC_NAME}"
             "&x-modelServer=4"
             "&x-imageTag=djl-inference:0.25.0-deepspeed0.11.0-cu118"
             f"&x-sdkVersion={SDK_VERSION}"
@@ -148,7 +153,7 @@ def test_capture_telemetry_decorator_tgi_success(self, mock_send_telemetry):
         args = mock_send_telemetry.call_args.args
         latency = str(args[5]).split("latency=")[1]
         expected_extra_str = (
-            f"{MOCK_FUNC_NAME}"
+            f"{MOCK_DEPLOY_FUNC_NAME}"
             "&x-modelServer=6"
             "&x-imageTag=huggingface-pytorch-inference:2.0.0-transformers4.28.1-cpu-py310-ubuntu20.04"
             f"&x-sdkVersion={SDK_VERSION}"
@@ -196,7 +201,7 @@ def test_capture_telemetry_decorator_handle_exception_success(self, mock_send_te
         args = mock_send_telemetry.call_args.args
         latency = str(args[5]).split("latency=")[1]
         expected_extra_str = (
-            f"{MOCK_FUNC_NAME}"
+            f"{MOCK_DEPLOY_FUNC_NAME}"
             "&x-modelServer=4"
             "&x-imageTag=djl-inference:0.25.0-deepspeed0.11.0-cu118"
             f"&x-sdkVersion={SDK_VERSION}"
@@ -243,7 +248,7 @@ def test_construct_url_with_failure_reason_and_extra_info(self):
             f"&x-failureType={mock_failure_type}"
             f"&x-extra={mock_extra_info}"
         )
-        self.assertEquals(ret_url, expected_base_url)
+        self.assertEqual(ret_url, expected_base_url)
 
     @patch("sagemaker.serve.utils.telemetry_logger._send_telemetry")
     def test_capture_telemetry_decorator_mlflow_success(self, mock_send_telemetry):
@@ -262,7 +267,7 @@ def test_capture_telemetry_decorator_mlflow_success(self, mock_send_telemetry):
         args = mock_send_telemetry.call_args.args
         latency = str(args[5]).split("latency=")[1]
         expected_extra_str = (
-            f"{MOCK_FUNC_NAME}"
+            f"{MOCK_DEPLOY_FUNC_NAME}"
             "&x-modelServer=1"
             "&x-imageTag=pytorch-inference:2.0.1-cpu-py310"
             f"&x-sdkVersion={SDK_VERSION}"
@@ -275,3 +280,66 @@ def test_capture_telemetry_decorator_mlflow_success(self, mock_send_telemetry):
         mock_send_telemetry.assert_called_once_with(
             "1", 3, MOCK_SESSION, None, None, expected_extra_str
         )
+
+    @patch("sagemaker.serve.utils.telemetry_logger._send_telemetry")
+    def test_capture_telemetry_decorator_optimize_with_default_configs(self, mock_send_telemetry):
+        mock_model_builder = ModelBuilderMock()
+        mock_model_builder.serve_settings.telemetry_opt_out = False
+        mock_model_builder.image_uri = None
+        mock_model_builder.mode = Mode.SAGEMAKER_ENDPOINT
+        mock_model_builder.model_server = ModelServer.TORCHSERVE
+        mock_model_builder.sagemaker_session.endpoint_arn = None
+
+        mock_model_builder.mock_optimize()
+
+        args = mock_send_telemetry.call_args.args
+        latency = str(args[5]).split("latency=")[1]
+        expected_extra_str = (
+            f"{MOCK_OPTIMIZE_FUNC_NAME}"
+            "&x-modelServer=1"
+            f"&x-sdkVersion={SDK_VERSION}"
+            f"&x-latency={latency}"
+        )
+
+        mock_send_telemetry.assert_called_once_with(
+            "1", 3, MOCK_SESSION, None, None, expected_extra_str
+        )
+
+    @patch("sagemaker.serve.utils.telemetry_logger._send_telemetry")
+    def test_capture_telemetry_decorator_optimize_with_custom_configs(self, mock_send_telemetry):
+        mock_model_builder = ModelBuilderMock()
+        mock_model_builder.serve_settings.telemetry_opt_out = False
+        mock_model_builder.image_uri = None
+        mock_model_builder.mode = Mode.SAGEMAKER_ENDPOINT
+        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_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,
+        )
+
+        args = mock_send_telemetry.call_args.args
+        latency = str(args[5]).split("latency=")[1]
+        expected_extra_str = (
+            f"{MOCK_OPTIMIZE_FUNC_NAME}"
+            "&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"
+            f"&x-latency={latency}"
+        )
+
+        mock_send_telemetry.assert_called_once_with(
+            "1", 3, MOCK_SESSION, None, None, expected_extra_str
+        )