Bugfixes from e2e testing. (#1670)

Author: cj-zhang

src/sagemaker/serve/builder/model_builder.py

@@ -250,7 +250,7 @@ class ModelBuilder(Triton, DJL, JumpStart, TGI, Transformers, TensorflowServing,
         default=None, metadata={"help": "Define sagemaker session for execution"}
     )
     name: Optional[str] = field(
-        default="model-name-" + uuid.uuid1().hex,
+        default_factory=lambda: "model-name-" + uuid.uuid1().hex,
         metadata={"help": "Define the model name"},
     )
     mode: Optional[Mode] = field(
@@ -1130,7 +1130,7 @@ def build(
     def _get_processing_unit(self):
         """Detects if the resource requirements are intended for a CPU or GPU instance."""
         # Assume custom orchestrator will be deployed as an endpoint to a CPU instance
-        if not self.resource_requirements:
+        if not self.resource_requirements or not self.resource_requirements.num_accelerators:
             return "cpu"
         for ic in self.modelbuilder_list or []:
             if ic.resource_requirements.num_accelerators > 0:
@@ -1171,10 +1171,10 @@ def _get_ic_resource_requirements(self, mb: ModelBuilder = None) -> ModelBuilder
 
     @_capture_telemetry("build_custom_orchestrator")
     def _get_smd_image_uri(self, processing_unit: str = None) -> str:
-        """Gets the SMD Inference URI.
+        """Gets the SMD Inference Image URI.
 
         Returns:
-            str: Pytorch DLC URI.
+            str: SMD Inference Image URI.
         """
         from sagemaker import image_uris
         import sys
@@ -1183,10 +1183,10 @@ def _get_smd_image_uri(self, processing_unit: str = None) -> str:
         from packaging.version import Version
 
         formatted_py_version = f"py{sys.version_info.major}{sys.version_info.minor}"
-        if Version(f"{sys.version_info.major}{sys.version_info.minor}") < Version("3.11.11"):
+        if Version(f"{sys.version_info.major}{sys.version_info.minor}") < Version("3.12"):
             raise ValueError(
                 f"Found Python version {formatted_py_version} but"
-                f"Custom orchestrator deployment requires Python version >= 3.11.11."
+                f"Custom orchestrator deployment requires Python version >= 3.12."
             )
 
         INSTANCE_TYPES = {"cpu": "ml.c5.xlarge", "gpu": "ml.g5.4xlarge"}
@@ -1957,7 +1957,7 @@ def deploy(
         ] = None,
         update_endpoint: Optional[bool] = False,
         custom_orchestrator_instance_type: str = None,
-        custom_orchestrator_initial_instance_count: int = 1,
+        custom_orchestrator_initial_instance_count: int = None,
         **kwargs,
     ) -> Union[Predictor, Transformer, List[Predictor]]:
         """Deploys the built Model.
@@ -2054,13 +2054,14 @@ def deploy(
             )
         if self._deployables.get("CustomOrchestrator", None):
             custom_orchestrator = self._deployables.get("CustomOrchestrator")
+            if not custom_orchestrator_instance_type and not instance_type:
+                logger.warning(
+                    "Deploying custom orchestrator as an endpoint but no instance type was "
+                    "set. Defaulting to `ml.c5.xlarge`."
+                )
+                custom_orchestrator_instance_type = "ml.c5.xlarge"
+                custom_orchestrator_initial_instance_count = 1
             if custom_orchestrator["Mode"] == "Endpoint":
-                if not custom_orchestrator_instance_type:
-                    logger.warning(
-                        "Deploying custom orchestrator as an endpoint but no instance type was "
-                        "set. Defaulting to `ml.c5.xlarge`."
-                    )
-                    custom_orchestrator_instance_type = "ml.c5.xlarge"
                 logger.info(
                     "Deploying custom orchestrator on instance type %s.",
                     custom_orchestrator_instance_type,
@@ -2073,13 +2074,18 @@ def deploy(
                     )
                 )
             elif custom_orchestrator["Mode"] == "InferenceComponent":
+                logger.info(
+                    "Deploying custom orchestrator as an inference component "
+                    f"to endpoint {endpoint_name}"
+                )
                 predictors.append(
                     self._deploy_for_ic(
                         ic_data=custom_orchestrator,
                         container_timeout_in_seconds=container_timeout_in_second,
                         instance_type=custom_orchestrator_instance_type or instance_type,
                         initial_instance_count=custom_orchestrator_initial_instance_count
                         or initial_instance_count,
+                        endpoint_name=endpoint_name,
                         **kwargs,
                     )
                 )

src/sagemaker/serve/model_server/smd/custom_execution_inference.py

@@ -67,6 +67,6 @@ async def handler(request):
     :return: outputs to be send back to client
     """
     if asyncio.iscoroutinefunction(custom_orchestrator.handle):
-        return await custom_orchestrator.handle(request)
+        return await custom_orchestrator.handle(request.body)
     else:
-        return custom_orchestrator.handle(request)
+        return custom_orchestrator.handle(request.body)

src/sagemaker/serve/spec/inference_base.py

@@ -24,11 +24,10 @@ def __init__(self):
     @property
     def client(self):
         """Boto3 SageMaker runtime client to use with custom orchestrator"""
-        if not hasattr(self, "_client"):
+        if not hasattr(self, "_client") or not self._client:
             from boto3 import Session
 
             self._client = Session().client("sagemaker-runtime")
-
         return self._client
 
     @abstractmethod