Fix formatting and add uts

Author: jiapinw

src/sagemaker/serve/builder/model_builder.py

@@ -294,7 +294,6 @@ def _auto_detect_container(self):
                 "Skipping auto detection as the image uri is provided %s",
                 self.image_uri,
             )
-            self._is_custom_image_uri = True
             return
 
         if self.model:

src/sagemaker/serve/utils/telemetry_logger.py

@@ -13,7 +13,6 @@
 """Placeholder docstring"""
 from __future__ import absolute_import
 import logging
-from typing import Optional
 from time import perf_counter
 
 import requests
@@ -69,6 +68,7 @@ def wrapper(self, *args, **kwargs):
                 f"&x-modelServer={MODEL_SERVER_TO_CODE[str(self.model_server)]}"
                 f"&x-imageTag={image_uri_tail}"
                 f"&x-sdkVersion={SDK_VERSION}"
+                f"&x-defaultImageUsage={_get_image_uri_option(self.image_uri, self._is_custom_image_uri)}"
             )
 
             if self.model_server == ModelServer.DJL_SERVING or self.model_server == ModelServer.TGI:
@@ -77,8 +77,6 @@ def wrapper(self, *args, **kwargs):
             if self.sagemaker_session and self.sagemaker_session.endpoint_arn:
                 extra += f"&x-endpointArn={self.sagemaker_session.endpoint_arn}"
 
-            extra += f"&x-defaultImageUsage={_get_image_uri_option(self.image_uri, self._is_custom_image_uri)}"
-
             start_timer = perf_counter()
             try:
                 response = func(self, *args, **kwargs)
@@ -95,10 +93,10 @@ def wrapper(self, *args, **kwargs):
                         extra,
                     )
             except (
-                    ModelBuilderException,
-                    exceptions.CapacityError,
-                    exceptions.UnexpectedStatusException,
-                    exceptions.AsyncInferenceError,
+                ModelBuilderException,
+                exceptions.CapacityError,
+                exceptions.UnexpectedStatusException,
+                exceptions.AsyncInferenceError,
             ) as e:
                 stop_timer = perf_counter()
                 elapsed = stop_timer - start_timer
@@ -126,12 +124,12 @@ def wrapper(self, *args, **kwargs):
 
 
 def _send_telemetry(
-        status: str,
-        mode: int,
-        session: Session,
-        failure_reason: str = None,
-        failure_type: str = None,
-        extra_info: str = None,
+    status: str,
+    mode: int,
+    session: Session,
+    failure_reason: str = None,
+    failure_type: str = None,
+    extra_info: str = None,
 ) -> None:
     """Make GET request to an empty object in S3 bucket"""
     try:
@@ -153,13 +151,13 @@ def _send_telemetry(
 
 
 def _construct_url(
-        accountId: str,
-        mode: str,
-        status: str,
-        failure_reason: str,
-        failure_type: str,
-        extra_info: str,
-        region: str,
+    accountId: str,
+    mode: str,
+    status: str,
+    failure_reason: str,
+    failure_type: str,
+    extra_info: str,
+    region: str,
 ) -> str:
     """Placeholder docstring"""
 

tests/unit/sagemaker/serve/builder/test_transformers_builder.py

@@ -16,6 +16,7 @@
 import unittest
 from sagemaker.serve.builder.model_builder import ModelBuilder
 from sagemaker.serve.mode.function_pointers import Mode
+from tests.unit.sagemaker.serve.constants import MOCK_VPC_CONFIG
 
 from sagemaker.serve.utils.predictors import TransformersLocalModePredictor
 
@@ -74,6 +75,7 @@ def test_build_deploy_for_transformers_local_container_and_remote_container(
             model=mock_model_id,
             schema_builder=mock_schema_builder,
             mode=Mode.LOCAL_CONTAINER,
+            vpc_config=MOCK_VPC_CONFIG,
         )
 
         builder._prepare_for_mode = MagicMock()
@@ -85,6 +87,7 @@ def test_build_deploy_for_transformers_local_container_and_remote_container(
         builder.modes[str(Mode.LOCAL_CONTAINER)] = MagicMock()
         predictor = model.deploy(model_data_download_timeout=1800)
 
+        assert model.vpc_config == MOCK_VPC_CONFIG
         assert builder.env_vars["MODEL_LOADING_TIMEOUT"] == "1800"
         assert isinstance(predictor, TransformersLocalModePredictor)
 

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.serve.utils.types import ImageUriOption
 from sagemaker.user_agent import SDK_VERSION
 
 MOCK_SESSION = Mock()
@@ -71,6 +72,7 @@ def test_capture_telemetry_decorator_djl_success(self, mock_send_telemetry):
         mock_model_builder = ModelBuilderMock()
         mock_model_builder.serve_settings.telemetry_opt_out = False
         mock_model_builder.image_uri = MOCK_DJL_CONTAINER
+        mock_model_builder._is_custom_image_uri = False
         mock_model_builder.model = MOCK_HUGGINGFACE_ID
         mock_model_builder.mode = Mode.LOCAL_CONTAINER
         mock_model_builder.model_server = ModelServer.DJL_SERVING
@@ -88,6 +90,37 @@ def test_capture_telemetry_decorator_djl_success(self, mock_send_telemetry):
             f"&x-modelName={MOCK_HUGGINGFACE_ID}"
             f"&x-endpointArn={MOCK_ENDPOINT_ARN}"
             f"&x-latency={latency}"
+            f"&x-defaultImageUsage={ImageUriOption.DEFAULT_IMAGE.value}"
+        )
+
+        mock_send_telemetry.assert_called_once_with(
+            "1", 2, MOCK_SESSION, None, None, expected_extra_str
+        )
+
+    @patch("sagemaker.serve.utils.telemetry_logger._send_telemetry")
+    def test_capture_telemetry_decorator_djl_success_with_custom_image(self, mock_send_telemetry):
+        mock_model_builder = ModelBuilderMock()
+        mock_model_builder.serve_settings.telemetry_opt_out = False
+        mock_model_builder.image_uri = MOCK_DJL_CONTAINER
+        mock_model_builder._is_custom_image_uri = True
+        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()
+
+        args = mock_send_telemetry.call_args.args
+        latency = str(args[5]).split("latency=")[1]
+        expected_extra_str = (
+            f"{MOCK_FUNC_NAME}"
+            "&x-modelServer=4"
+            "&x-imageTag=djl-inference:0.25.0-deepspeed0.11.0-cu118"
+            f"&x-sdkVersion={SDK_VERSION}"
+            f"&x-modelName={MOCK_HUGGINGFACE_ID}"
+            f"&x-endpointArn={MOCK_ENDPOINT_ARN}"
+            f"&x-latency={latency}"
+            f"&x-defaultImageUsage={ImageUriOption.CUSTOM_1P_IMAGE.value}"
         )
 
         mock_send_telemetry.assert_called_once_with(
@@ -99,6 +132,7 @@ def test_capture_telemetry_decorator_tgi_success(self, mock_send_telemetry):
         mock_model_builder = ModelBuilderMock()
         mock_model_builder.serve_settings.telemetry_opt_out = False
         mock_model_builder.image_uri = MOCK_TGI_CONTAINER
+        mock_model_builder._is_custom_image_uri = False
         mock_model_builder.model = MOCK_HUGGINGFACE_ID
         mock_model_builder.mode = Mode.LOCAL_CONTAINER
         mock_model_builder.model_server = ModelServer.TGI
@@ -116,6 +150,7 @@ def test_capture_telemetry_decorator_tgi_success(self, mock_send_telemetry):
             f"&x-modelName={MOCK_HUGGINGFACE_ID}"
             f"&x-endpointArn={MOCK_ENDPOINT_ARN}"
             f"&x-latency={latency}"
+            f"&x-defaultImageUsage={ImageUriOption.DEFAULT_IMAGE.value}"
         )
 
         mock_send_telemetry.assert_called_once_with(
@@ -139,6 +174,7 @@ def test_capture_telemetry_decorator_handle_exception_success(self, mock_send_te
         mock_model_builder = ModelBuilderMock()
         mock_model_builder.serve_settings.telemetry_opt_out = False
         mock_model_builder.image_uri = MOCK_DJL_CONTAINER
+        mock_model_builder._is_custom_image_uri = False
         mock_model_builder.model = MOCK_HUGGINGFACE_ID
         mock_model_builder.mode = Mode.LOCAL_CONTAINER
         mock_model_builder.model_server = ModelServer.DJL_SERVING
@@ -161,6 +197,7 @@ def test_capture_telemetry_decorator_handle_exception_success(self, mock_send_te
             f"&x-modelName={MOCK_HUGGINGFACE_ID}"
             f"&x-endpointArn={MOCK_ENDPOINT_ARN}"
             f"&x-latency={latency}"
+            f"&x-defaultImageUsage={ImageUriOption.DEFAULT_IMAGE.value}"
         )
 
         mock_send_telemetry.assert_called_once_with(