Emit additional telemetry metrics

Author: Jonathan Makunga

.gitignore

@@ -35,3 +35,4 @@ env/
 tests/data/**/_repack_model.py
 tests/data/experiment/sagemaker-dev-1.0.tar.gz
 src/sagemaker/serve/tmp_workspace
+src/sagemaker/image_uri_config/pysdk_version.json

setup.py

@@ -32,6 +32,21 @@ def read_version():
     return read("VERSION").strip()
 
 
+def pysdk_version():
+    """Persists Sagemaker Python SDK Version in Config"""
+    content = '{"version": "' + read_version() + '"}'
+    with open(
+        os.path.join(
+            os.path.dirname(__file__), "src", "sagemaker", "image_uri_config", "pysdk_version.json"
+        ),
+        "w",
+    ) as v:
+        v.write(content)
+
+
+pysdk_version()
+
+
 def read_requirements(filename):
     """Reads requirements file which lists package dependencies.
 

src/sagemaker/serve/utils/telemetry_logger.py

@@ -15,10 +15,11 @@
 import logging
 import requests
 
-from sagemaker import Session
+from sagemaker import Session, exceptions
 from sagemaker.serve.mode.function_pointers import Mode
 from sagemaker.serve.utils.exceptions import ModelBuilderException
 from sagemaker.serve.utils.types import ModelServer
+from sagemaker.utils import pysdk_version
 
 logger = logging.getLogger(__name__)
 
@@ -63,11 +64,15 @@ def wrapper(self, *args, **kwargs):
                 f"{func_name}"
                 f"&x-modelServer={MODEL_SERVER_TO_CODE[str(self.model_server)]}"
                 f"&x-imageTag={image_uri_tail}"
+                f"&x-pySdkVersion={pysdk_version()}"
             )
 
             if self.model_server == ModelServer.DJL_SERVING or self.model_server == ModelServer.TGI:
                 extra += f"&x-modelName={self.model}"
 
+            if self.sagemaker_session.endpoint_arn:
+                extra += f"&x-endpointArn={self.sagemaker_session.endpoint_arn}"
+
             try:
                 response = func(self, *args, **kwargs)
                 if not self.serve_settings.telemetry_opt_out:
@@ -79,7 +84,11 @@ def wrapper(self, *args, **kwargs):
                         None,
                         extra,
                     )
-            except ModelBuilderException as e:
+            except (
+                ModelBuilderException,
+                exceptions.CapacityError,
+                exceptions.UnexpectedStatusException,
+            ) as e:
                 if not self.serve_settings.telemetry_opt_out:
                     _send_telemetry(
                         "0",

src/sagemaker/session.py

@@ -243,6 +243,7 @@ def __init__(
         # sagemaker_config is validated and initialized inside :func:`_initialize`,
         # so if default_bucket is None and the sagemaker_config has a default S3 bucket configured,
         # _default_bucket_name_override will be set again inside :func:`_initialize`.
+        self.endpoint_arn = None
         self._default_bucket = None
         self._default_bucket_name_override = default_bucket
         # this may also be set again inside :func:`_initialize` if it is None
@@ -5054,6 +5055,7 @@ def wait_for_endpoint(self, endpoint, poll=DEFAULT_EP_POLL, live_logging=False):
                 poll=EP_LOGGER_POLL,
             )
         status = desc["EndpointStatus"]
+        self.endpoint_arn = desc["EndpointArn"]
 
         if status != "InService":
             reason = desc.get("FailureReason", None)

src/sagemaker/utils.py

@@ -1489,3 +1489,10 @@ def format_tags(tags: Tags) -> List[TagsDict]:
         return [{"Key": str(k), "Value": str(v)} for k, v in tags.items()]
 
     return tags
+
+
+def pysdk_version() -> str:
+    """Returns the current Sagemaker Python SDK Version"""
+    v_path = os.path.join(os.path.dirname(__file__), "image_uri_config", "pysdk_version.json")
+    with open(v_path) as v:
+        return json.load(v).get("version")

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

@@ -20,6 +20,7 @@
     _construct_url,
 )
 from sagemaker.serve.utils.exceptions import ModelBuilderException, LocalModelOutOfMemoryException
+from sagemaker.utils import pysdk_version
 
 MOCK_SESSION = Mock()
 MOCK_FUNC_NAME = "Mock.deploy"
@@ -32,6 +33,10 @@
 )
 MOCK_HUGGINGFACE_ID = "meta-llama/Llama-2-7b-hf"
 MOCK_EXCEPTION = LocalModelOutOfMemoryException("mock raise ex")
+MOCK_ENDPOINT_ARN = (
+    "arn:aws:sagemaker:us-west-2:123456789012:endpoint/huggingface-pytorch-tgi-inference-2024-02-06"
+    "-04-06-23-819"
+)
 
 
 class ModelBuilderMock:
@@ -72,14 +77,17 @@ def test_capture_telemetry_decorator_djl_success(self, mock_send_telemetry):
         mock_model_builder.model = MOCK_HUGGINGFACE_ID
         mock_model_builder.mode = Mode.LOCAL_CONTAINER
         mock_model_builder.model_server = ModelServer.DJL_SERVING
+        mock_model_builder.sagemaker_session.endpoint_arn = MOCK_ENDPOINT_ARN
 
         mock_model_builder.mock_deploy()
 
         expected_extra_str = (
             f"{MOCK_FUNC_NAME}"
             "&x-modelServer=4"
             "&x-imageTag=djl-inference:0.25.0-deepspeed0.11.0-cu118"
+            f"&x-pySdkVersion={pysdk_version()}"
             f"&x-modelName={MOCK_HUGGINGFACE_ID}"
+            f"&x-endpointArn={MOCK_ENDPOINT_ARN}"
         )
         mock_send_telemetry.assert_called_once_with(
             "1", 2, MOCK_SESSION, None, None, expected_extra_str
@@ -93,14 +101,17 @@ def test_capture_telemetry_decorator_tgi_success(self, mock_send_telemetry):
         mock_model_builder.model = MOCK_HUGGINGFACE_ID
         mock_model_builder.mode = Mode.LOCAL_CONTAINER
         mock_model_builder.model_server = ModelServer.TGI
+        mock_model_builder.sagemaker_session.endpoint_arn = MOCK_ENDPOINT_ARN
 
         mock_model_builder.mock_deploy()
 
         expected_extra_str = (
             f"{MOCK_FUNC_NAME}"
             "&x-modelServer=6"
             "&x-imageTag=huggingface-pytorch-inference:2.0.0-transformers4.28.1-cpu-py310-ubuntu20.04"
+            f"&x-pySdkVersion={pysdk_version()}"
             f"&x-modelName={MOCK_HUGGINGFACE_ID}"
+            f"&x-endpointArn={MOCK_ENDPOINT_ARN}"
         )
         mock_send_telemetry.assert_called_once_with(
             "1", 2, MOCK_SESSION, None, None, expected_extra_str
@@ -126,6 +137,7 @@ def test_capture_telemetry_decorator_handle_exception_success(self, mock_send_te
         mock_model_builder.model = MOCK_HUGGINGFACE_ID
         mock_model_builder.mode = Mode.LOCAL_CONTAINER
         mock_model_builder.model_server = ModelServer.DJL_SERVING
+        mock_model_builder.sagemaker_session.endpoint_arn = MOCK_ENDPOINT_ARN
 
         mock_exception = Mock()
         mock_exception_obj = MOCK_EXCEPTION
@@ -138,7 +150,9 @@ def test_capture_telemetry_decorator_handle_exception_success(self, mock_send_te
             f"{MOCK_FUNC_NAME}"
             "&x-modelServer=4"
             "&x-imageTag=djl-inference:0.25.0-deepspeed0.11.0-cu118"
+            f"&x-pySdkVersion={pysdk_version()}"
             f"&x-modelName={MOCK_HUGGINGFACE_ID}"
+            f"&x-endpointArn={MOCK_ENDPOINT_ARN}"
         )
         mock_send_telemetry.assert_called_once_with(
             "0",

tests/unit/test_utils.py

@@ -42,6 +42,7 @@
     resolve_nested_dict_value_from_config,
     update_list_of_dicts_with_values_from_config,
     volume_size_supported,
+    pysdk_version,
 )
 from tests.unit.sagemaker.workflow.helpers import CustomStep
 from sagemaker.workflow.parameters import ParameterString, ParameterInteger
@@ -1748,3 +1749,6 @@ def test_instance_family_from_full_instance_type(self):
 
         for instance_type, family in instance_type_to_family_test_dict.items():
             self.assertEqual(family, get_instance_type_family(instance_type))
+
+    def test_pysdk_version(self):
+        self.assertIsNotNone(pysdk_version())