fix: Address SA feedback regarding deployment straight to Endpoint Mode - Galactus (#1405)

Author: gwang111

src/sagemaker/serve/builder/djl_builder.py

@@ -36,8 +36,14 @@
     _set_serve_properties,
     _get_admissible_tensor_parallel_degrees,
     _get_admissible_dtypes,
+    _get_default_tensor_parallel_degree,
+)
+from sagemaker.serve.utils.local_hardware import (
+    _get_nb_instance,
+    _get_ram_usage_mb,
+    _get_gpu_info,
+    _get_gpu_info_fallback,
 )
-from sagemaker.serve.utils.local_hardware import _get_nb_instance, _get_ram_usage_mb
 from sagemaker.serve.model_server.djl_serving.prepare import (
     prepare_for_djl_serving,
     _create_dir_structure,
@@ -164,13 +170,6 @@ def _create_djl_model(self) -> Type[Model]:
     @_capture_telemetry("djl.deploy")
     def _djl_model_builder_deploy_wrapper(self, *args, **kwargs) -> Type[PredictorBase]:
         """Placeholder docstring"""
-        prepare_for_djl_serving(
-            model_path=self.model_path,
-            model=self.pysdk_model,
-            dependencies=self.dependencies,
-            overwrite_props_from_file=self.overwrite_props_from_file,
-        )
-
         timeout = kwargs.get("model_data_download_timeout")
         if timeout:
             self.env_vars.update({"MODEL_LOADING_TIMEOUT": str(timeout)})
@@ -192,6 +191,34 @@ def _djl_model_builder_deploy_wrapper(self, *args, **kwargs) -> Type[PredictorBa
             else:
                 raise ValueError("Mode %s is not supported!" % overwrite_mode)
 
+        manual_set_props = None
+        if self.mode == Mode.SAGEMAKER_ENDPOINT:
+            if self.nb_instance_type and "instance_type" not in kwargs:
+                kwargs.update({"instance_type": self.nb_instance_type})
+            elif not self.nb_instance_type and "instance_type" not in kwargs:
+                raise ValueError(
+                    "Instance type must be provided when deploying " "to SageMaker Endpoint mode."
+                )
+            else:
+                try:
+                    tot_gpus = _get_gpu_info(kwargs.get("instance_type"), self.sagemaker_session)
+                except Exception:  # pylint: disable=W0703
+                    tot_gpus = _get_gpu_info_fallback(kwargs.get("instance_type"))
+                default_tensor_parallel_degree = _get_default_tensor_parallel_degree(
+                    self.hf_model_config, tot_gpus
+                )
+                manual_set_props = {
+                    "option.tensor_parallel_degree": str(default_tensor_parallel_degree) + "\n"
+                }
+
+        prepare_for_djl_serving(
+            model_path=self.model_path,
+            model=self.pysdk_model,
+            dependencies=self.dependencies,
+            overwrite_props_from_file=self.overwrite_props_from_file,
+            manual_set_props=manual_set_props,
+        )
+
         serializer = self.schema_builder.input_serializer
         deserializer = self.schema_builder._output_deserializer
         if self.mode == Mode.LOCAL_CONTAINER:
@@ -237,8 +264,6 @@ def _djl_model_builder_deploy_wrapper(self, *args, **kwargs) -> Type[PredictorBa
 
         if "endpoint_logging" not in kwargs:
             kwargs["endpoint_logging"] = True
-        if self.nb_instance_type and "instance_type" not in kwargs:
-            kwargs.update({"instance_type": self.nb_instance_type})
 
         predictor = self._original_deploy(*args, **kwargs)
 
@@ -252,6 +277,7 @@ def _build_for_hf_djl(self):
         """Placeholder docstring"""
         self.overwrite_props_from_file = True
         self.nb_instance_type = _get_nb_instance()
+
         _create_dir_structure(self.model_path)
         self.engine, self.hf_model_config = _auto_detect_engine(
             self.model, self.env_vars.get("HUGGING_FACE_HUB_TOKEN")

src/sagemaker/serve/builder/tgi_builder.py

@@ -32,13 +32,21 @@
     _pretty_print_results_tgi,
 )
 from sagemaker.djl_inference.model import _get_model_config_properties_from_hf
-from sagemaker.serve.model_server.djl_serving.utils import _get_admissible_tensor_parallel_degrees
+from sagemaker.serve.model_server.djl_serving.utils import (
+    _get_admissible_tensor_parallel_degrees,
+    _get_default_tensor_parallel_degree,
+)
 from sagemaker.serve.model_server.tgi.utils import (
     _get_default_tgi_configurations,
     _get_admissible_dtypes,
 )
 from sagemaker.huggingface import HuggingFaceModel, get_huggingface_llm_image_uri
-from sagemaker.serve.utils.local_hardware import _get_nb_instance, _get_ram_usage_mb
+from sagemaker.serve.utils.local_hardware import (
+    _get_nb_instance,
+    _get_ram_usage_mb,
+    _get_gpu_info,
+    _get_gpu_info_fallback,
+)
 from sagemaker.serve.model_server.tgi.prepare import _create_dir_structure
 from sagemaker.serve.utils.predictors import TgiLocalModePredictor
 from sagemaker.serve.utils.types import ModelServer
@@ -202,8 +210,26 @@ def _tgi_model_builder_deploy_wrapper(self, *args, **kwargs) -> Type[PredictorBa
 
         if "endpoint_logging" not in kwargs:
             kwargs["endpoint_logging"] = True
+
         if self.nb_instance_type and "instance_type" not in kwargs:
             kwargs.update({"instance_type": self.nb_instance_type})
+        elif not self.nb_instance_type and "instance_type" not in kwargs:
+            raise ValueError(
+                "Instance type must be provided when deploying " "to SageMaker Endpoint mode."
+            )
+        else:
+            try:
+                tot_gpus = _get_gpu_info(kwargs.get("instance_type"), self.sagemaker_session)
+            except Exception:  # pylint: disable=W0703
+                tot_gpus = _get_gpu_info_fallback(kwargs.get("instance_type"))
+            default_num_shard = _get_default_tensor_parallel_degree(self.hf_model_config, tot_gpus)
+            self.pysdk_model.env.update(
+                {
+                    "NUM_SHARD": str(default_num_shard),
+                    "SHARDED": "true" if default_num_shard > 1 else "false",
+                }
+            )
+
         if "initial_instance_count" not in kwargs:
             kwargs.update({"initial_instance_count": 1})
 
@@ -218,6 +244,7 @@ def _tgi_model_builder_deploy_wrapper(self, *args, **kwargs) -> Type[PredictorBa
     def _build_for_hf_tgi(self):
         """Placeholder docstring"""
         self.nb_instance_type = _get_nb_instance()
+
         _create_dir_structure(self.model_path)
         if not hasattr(self, "pysdk_model"):
             self.env_vars.update({"HF_MODEL_ID": self.model})

src/sagemaker/serve/mode/local_container_mode.py

@@ -190,5 +190,5 @@ def _pull_image(self, image: str):
         try:
             logger.info("Pulling image %s from repository...", image)
             self.client.images.pull(image)
-        except docker.errors.NotFound:
-            logger.warning("Could not find remote image to pull")
+        except docker.errors.NotFound as e:
+            raise ValueError("Could not find remote image to pull") from e

src/sagemaker/serve/model_server/djl_serving/prepare.py

@@ -1,3 +1,15 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
 """Prepare DjlModel for Deployment"""
 
 from __future__ import absolute_import
@@ -55,7 +67,9 @@ def _copy_jumpstart_artifacts(model_data: str, js_id: str, code_dir: Path):
     return (existing_properties, hf_model_config, True)
 
 
-def _generate_properties_file(model: DJLModel, code_dir: Path, overwrite_props_from_file: bool):
+def _generate_properties_file(
+    model: DJLModel, code_dir: Path, overwrite_props_from_file: bool, manual_set_props: dict
+):
     """Placeholder Docstring"""
     if _has_serving_properties_file(code_dir):
         existing_properties = _read_existing_serving_properties(code_dir)
@@ -67,6 +81,13 @@ def _generate_properties_file(model: DJLModel, code_dir: Path, overwrite_props_f
 
     with open(serving_properties_file, mode="w+") as file:
         covered_keys = set()
+
+        if manual_set_props:
+            for key, value in manual_set_props.items():
+                logger.info(_SETTING_PROPERTY_STMT, key, value.strip())
+                covered_keys.add(key)
+                file.write(f"{key}={value}")
+
         for key, value in serving_properties_dict.items():
             if not overwrite_props_from_file:
                 logger.info(_SETTING_PROPERTY_STMT, key, value)
@@ -129,6 +150,7 @@ def prepare_for_djl_serving(
     shared_libs: List[str] = None,
     dependencies: str = None,
     overwrite_props_from_file: bool = True,
+    manual_set_props: dict = None,
 ):
     """Prepare serving when a HF model id is given
 
@@ -149,7 +171,7 @@ def prepare_for_djl_serving(
 
     _copy_inference_script(code_dir)
 
-    _generate_properties_file(model, code_dir, overwrite_props_from_file)
+    _generate_properties_file(model, code_dir, overwrite_props_from_file, manual_set_props)
 
 
 def prepare_djl_js_resources(

src/sagemaker/serve/model_server/djl_serving/utils.py

@@ -60,9 +60,11 @@ def _auto_detect_engine(model_id: str, hf_hub_token: str) -> tuple:
     return (engine, hf_model_config)
 
 
-def _get_default_tensor_parallel_degree(hf_model_config: dict) -> int:
+def _get_default_tensor_parallel_degree(hf_model_config: dict, gpu_count: int = None) -> int:
     """Placeholder docstring"""
     available_gpus = _get_available_gpus()
+    if not available_gpus and not gpu_count:
+        return None
 
     attention_heads = None
     for variant in ATTENTION_HEAD_NAME_VARIENTS:
@@ -73,7 +75,8 @@ def _get_default_tensor_parallel_degree(hf_model_config: dict) -> int:
     if not attention_heads:
         return 1
 
-    for i in (n + 1 for n in reversed(range(len(available_gpus)))):
+    tot_gpus = len(available_gpus) if available_gpus else gpu_count
+    for i in (n + 1 for n in reversed(range(tot_gpus))):
         if attention_heads % i == 0:
             logger.info(
                 "Max GPU parallelism of %s is allowed. Total attention heads %s", i, attention_heads

src/sagemaker/serve/model_server/tgi/utils.py

@@ -23,10 +23,19 @@ def _get_default_tgi_configurations(
         schema_builder.sample_input, schema_builder.sample_output
     )
 
+    if default_num_shard:
+        return (
+            {
+                "SHARDED": "true" if default_num_shard > 1 else "false",
+                "NUM_SHARD": str(default_num_shard),
+                "DTYPE": _get_default_dtype(),
+            },
+            default_max_new_tokens,
+        )
     return (
         {
-            "SHARDED": "true" if default_num_shard > 1 else "false",
-            "NUM_SHARD": str(default_num_shard),
+            "SHARDED": None,
+            "NUM_SHARD": None,
             "DTYPE": _get_default_dtype(),
         },
         default_max_new_tokens,

src/sagemaker/serve/utils/local_hardware.py

@@ -1,3 +1,15 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
 """Utilites for identifying and analyzing local gpu hardware"""
 from __future__ import absolute_import
 
@@ -10,6 +22,8 @@
 from pathlib import Path
 import psutil
 
+from sagemaker import Session
+
 logger = logging.getLogger(__name__)
 
 # key = vCPUs
@@ -54,6 +68,34 @@
 }
 
 
+fallback_gpu_resource_mapping = {
+    "ml.p5.48xlarge": 8,
+    "ml.p4d.24xlarge": 8,
+    "ml.p4de.24xlarge": 8,
+    "ml.p3.2xlarge": 1,
+    "ml.p3.8xlarge": 4,
+    "ml.p3.16xlarge": 8,
+    "ml.p3dn.24xlarge": 8,
+    "ml.p2.xlarge": 1,
+    "ml.p2.8xlarge": 8,
+    "ml.p2.16xlarge": 16,
+    "ml.g4dn.xlarge": 1,
+    "ml.g4dn.2xlarge": 1,
+    "ml.g4dn.4xlarge": 1,
+    "ml.g4dn.8xlarge": 1,
+    "ml.g4dn.16xlarge": 1,
+    "ml.g4dn.12xlarge": 4,
+    "ml.g5n.xlarge": 1,
+    "ml.g5.2xlarge": 1,
+    "ml.g5.4xlarge": 1,
+    "ml.g5.8xlarge": 1,
+    "ml.g5.16xlarge": 1,
+    "ml.g5.12xlarge": 4,
+    "ml.g5.24xlarge": 4,
+    "ml.g5.48xlarge": 8,
+}
+
+
 def _get_available_gpus(log=True):
     """Detect the GPUs available on the device and their available resources"""
     try:
@@ -63,16 +105,24 @@ def _get_available_gpus(log=True):
 
         if log:
             logger.info("CUDA enabled hardware on the device: %s", gpu_info)
-
         return gpu_info
-    except Exception as e:
+    except Exception as e:  # pylint: disable=W0703
         # for nvidia-smi to run, a cuda driver must be present
-        raise ValueError("CUDA is not enabled on your device. %s" % str(e))
+        logger.warning(
+            "CUDA is not enabled on your device. %s. "
+            "Please run ModelBuilder on CUDA enabled hardware "
+            "to deploy locally.",
+            str(e),
+        )
+        return None
 
 
 def _get_nb_instance():
     """Placeholder docstring"""
     gpu_info = _get_available_gpus(False)
+    if not gpu_info:
+        return None
+
     gpu_name, gpu_mem = gpu_info[0].split(", ")
     cpu_count = multiprocessing.cpu_count()
 
@@ -156,3 +206,29 @@ def _check_docker_disk_usage():
             docker_path,
             str(e),
         )
+
+
+def _get_gpu_info(instance_type: str, session: Session) -> int:
+    """Get GPU info for the provided instance"""
+    ec2_client = session.boto_session.client("ec2")
+
+    split_instance = instance_type.split(".")
+    split_instance.pop(0)
+
+    ec2_instance = ".".join(split_instance)
+
+    instance_info = ec2_client.describe_instance_types(InstanceTypes=[ec2_instance])
+
+    gpus_info = instance_info.get("InstanceTypes")[0].get("GpuInfo")
+
+    if gpus_info:
+        return gpus_info.get("Gpus")[0].get("Count")
+    raise ValueError("Provided instance_type is not GPU enabled.")
+
+
+def _get_gpu_info_fallback(instance_type: str) -> int:
+    """Get GPU info for the provided instance fallback"""
+    available_gpus = fallback_gpu_resource_mapping.get(instance_type)
+    if not available_gpus:
+        raise ValueError("Provided instance_type is not GPU enabled.")
+    return available_gpus

src/sagemaker/serve/utils/predictors.py

@@ -137,16 +137,18 @@ def __init__(
 
     def predict(self, data):
         """Placeholder docstring"""
-        return self.deserializer.deserialize(
-            io.BytesIO(
-                self._mode_obj._invoke_tgi_serving(
-                    self.serializer.serialize(data),
-                    self.content_type,
-                    self.deserializer.ACCEPT[0],
-                )
-            ),
-            self.content_type,
-        )
+        return [
+            self.deserializer.deserialize(
+                io.BytesIO(
+                    self._mode_obj._invoke_tgi_serving(
+                        self.serializer.serialize(data),
+                        self.content_type,
+                        self.deserializer.ACCEPT[0],
+                    )
+                ),
+                self.content_type,
+            )
+        ]
 
     @property
     def content_type(self):