Initial commit

Author: jiapinw

src/sagemaker/serve/builder/djl_builder.py

@@ -83,6 +83,7 @@ def __init__(self):
         self.mode = None
         self.model_server = None
         self.image_uri = None
+        self._is_custom_image_uri = False
         self.image_config = None
         self.vpc_config = None
         self._original_deploy = None

src/sagemaker/serve/builder/jumpstart_builder.py

@@ -57,6 +57,8 @@ def __init__(self):
         self.mode = None
         self.model_server = None
         self.image_uri = None
+        self._is_custom_image_uri = False
+        self.vpc_config = None
         self._original_deploy = None
         self.secret_key = None
         self.js_model_config = None
@@ -94,7 +96,7 @@ def _is_jumpstart_model_id(self) -> bool:
 
     def _create_pre_trained_js_model(self) -> Type[Model]:
         """Placeholder docstring"""
-        pysdk_model = JumpStartModel(self.model)
+        pysdk_model = JumpStartModel(self.model, vpc_config=self.vpc_config)
         pysdk_model.sagemaker_session = self.sagemaker_session
 
         self._original_deploy = pysdk_model.deploy

src/sagemaker/serve/builder/model_builder.py

@@ -294,6 +294,7 @@ 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:
@@ -605,6 +606,8 @@ def build(
 
         self.serve_settings = self._get_serve_setting()
 
+        self._is_custom_image_uri = self.image_uri is None
+
         if isinstance(self.model, str):
             if self._is_jumpstart_model_id():
                 return self._build_for_jumpstart()

src/sagemaker/serve/builder/tgi_builder.py

@@ -76,6 +76,7 @@ def __init__(self):
         self.mode = None
         self.model_server = None
         self.image_uri = None
+        self._is_custom_image_uri = False
         self.image_config = None
         self.vpc_config = None
         self._original_deploy = None

src/sagemaker/serve/builder/transformers_builder.py

@@ -56,6 +56,8 @@ def __init__(self):
         self.mode = None
         self.model_server = None
         self.image_uri = None
+        self._is_custom_image_uri = False
+        self.vpc_config = None
         self._original_deploy = None
         self.hf_model_config = None
         self._default_data_type = None
@@ -111,6 +113,7 @@ def _create_transformers_model(self) -> Type[Model]:
                 py_version=self.py_version,
                 transformers_version=base_hf_version,
                 pytorch_version=self.pytorch_version,
+                vpc_config=self.vpc_config,
             )
         elif "keras" in hf_model_md.get("tags") or "tensorflow" in hf_model_md.get("tags"):
             self.tensorflow_version = self._get_supported_version(
@@ -126,6 +129,7 @@ def _create_transformers_model(self) -> Type[Model]:
                 py_version=self.py_version,
                 transformers_version=base_hf_version,
                 tensorflow_version=self.tensorflow_version,
+                vpc_config=self.vpc_config,
             )
 
         if self.mode == Mode.LOCAL_CONTAINER:

src/sagemaker/serve/utils/telemetry_logger.py

@@ -13,14 +13,16 @@
 """Placeholder docstring"""
 from __future__ import absolute_import
 import logging
+from typing import Optional
 from time import perf_counter
 
 import requests
 
 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.serve.utils.types import ModelServer, ImageUriOption
+from sagemaker.serve.validations.check_image_uri import is_1p_image_uri
 from sagemaker.user_agent import SDK_VERSION
 
 logger = logging.getLogger(__name__)
@@ -75,6 +77,8 @@ 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)
@@ -91,10 +95,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
@@ -122,12 +126,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:
@@ -149,13 +153,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"""
 
@@ -201,3 +205,22 @@ def _get_region_or_default(session):
         return session.boto_session.region_name
     except Exception:  # pylint: disable=W0703
         return "us-west-2"
+
+
+def _get_image_uri_option(image_uri: str, is_custom_image: bool) -> int:
+    """Detect whether default values are used for ModelBuilder
+
+    Args:
+        image_uri (str): Image uri used by ModelBuilder.
+        is_custom_image: (bool): Boolean indicating whether customer provides with custom image.
+    Returns:
+        bool: Integer code of image option types.
+    """
+
+    if not is_custom_image:
+        return ImageUriOption.DEFAULT_IMAGE.value
+
+    if is_1p_image_uri(image_uri):
+        return ImageUriOption.CUSTOM_1P_IMAGE.value
+
+    return ImageUriOption.CUSTOM_IMAGE.value

src/sagemaker/serve/utils/types.py

@@ -43,3 +43,15 @@ def __str__(self) -> str:
     INFERENTIA_1 = 3
     INFERENTIA_2 = 4
     GRAVITON = 5
+
+
+class ImageUriOption(Enum):
+    "Enum type for image uri options"
+
+    def __str__(self) -> str:
+        """Convert enum to string"""
+        return str(self.name)
+
+    CUSTOM_IMAGE = 1
+    CUSTOM_1P_IMAGE = 2
+    DEFAULT_IMAGE = 3