fix formatting and msging

gwang111 · gwang111 · commit 3d043848a3f3 · 2024-11-15T17:19:40.000Z
diff --git a/src/sagemaker/model.py b/src/sagemaker/model.py
@@ -1601,8 +1601,10 @@ def deploy(
                 self._base_name = "-".join((self._base_name, compiled_model_suffix))
 
         if self._is_sharded_model and endpoint_type != EndpointType.INFERENCE_COMPONENT_BASED:
-            logging.warning("Forcing INFERENCE_COMPONENT_BASED endpoint for sharded model. ADVISORY - "
-                            "Use INFERENCE_COMPONENT_BASED endpoints over MODEL_BASED endpoints.")
+            logging.warning(
+                "Forcing INFERENCE_COMPONENT_BASED endpoint for sharded model. ADVISORY - "
+                "Use INFERENCE_COMPONENT_BASED endpoints over MODEL_BASED endpoints."
+            )
             endpoint_type = EndpointType.INFERENCE_COMPONENT_BASED
 
         # Support multiple models on same endpoint
diff --git a/src/sagemaker/serve/builder/model_builder.py b/src/sagemaker/serve/builder/model_builder.py
@@ -104,7 +104,7 @@
     get_huggingface_model_metadata,
     download_huggingface_model_metadata,
 )
-from sagemaker.serve.validations.optimization import validate_optimization_configuration
+from sagemaker.serve.validations.optimization import _validate_optimization_configuration
 
 logger = logging.getLogger(__name__)
 
@@ -1161,15 +1161,6 @@ def optimize(
             Model: A deployable ``Model`` object.
         """
 
-        # TODO: ideally these dictionaries need to be sagemaker_core shapes
-        validate_optimization_configuration(
-            instance_type=instance_type,
-            quantization_config=quantization_config,
-            compilation_config=compilation_config,
-            sharding_config=sharding_config,
-            speculative_decoding_config=speculative_decoding_config,
-        )
-
         # need to get telemetry_opt_out info before telemetry decorator is called
         self.serve_settings = self._get_serve_setting()
 
@@ -1243,6 +1234,17 @@ def _model_builder_optimize_wrapper(
         Returns:
             Model: A deployable ``Model`` object.
         """
+
+        # TODO: ideally these dictionaries need to be sagemaker_core shapes
+        # TODO: for organization, abstract all validation behind this fn
+        _validate_optimization_configuration(
+            instance_type=instance_type,
+            quantization_config=quantization_config,
+            compilation_config=compilation_config,
+            sharding_config=sharding_config,
+            speculative_decoding_config=speculative_decoding_config,
+        )
+
         self.is_compiled = compilation_config is not None
         self.is_quantized = quantization_config is not None
         self.speculative_decoding_draft_model_source = _extract_speculative_draft_model_provider(
@@ -1255,11 +1257,29 @@ def _model_builder_optimize_wrapper(
         if quantization_config and compilation_config:
             raise ValueError("Quantization config and compilation config are mutually exclusive.")
 
-        if sharding_config and (quantization_config or compilation_config or speculative_decoding_config):
-            raise ValueError("Sharding config is mutually exclusive and cannot be combined with any other optimization.")
+        if sharding_config and (
+            quantization_config or compilation_config or speculative_decoding_config
+        ):
+            raise ValueError(
+                (
+                    "Sharding config is mutually exclusive "
+                    "and cannot be combined with any other optimization."
+                )
+            )
 
-        if sharding_config and ((env_vars and "OPTION_TENSOR_PARALLEL_DEGREE" not in env_vars) or (sharding_config.get("OverrideEnvironment") and "OPTION_TENSOR_PARALLEL_DEGREE" not in sharding_config["OverrideEnvironment"])):
-            raise ValueError("OPTION_TENSOR_PARALLEL_DEGREE is required environment variable with Sharding config.")
+        if sharding_config and (
+            (env_vars and "OPTION_TENSOR_PARALLEL_DEGREE" not in env_vars)
+            or (
+                sharding_config.get("OverrideEnvironment")
+                and "OPTION_TENSOR_PARALLEL_DEGREE" not in sharding_config["OverrideEnvironment"]
+            )
+        ):
+            raise ValueError(
+                (
+                    "OPTION_TENSOR_PARALLEL_DEGREE is required "
+                    "environment variable with Sharding config."
+                )
+            )
 
         self.sagemaker_session = sagemaker_session or self.sagemaker_session or Session()
         self.instance_type = instance_type or self.instance_type
diff --git a/src/sagemaker/serve/utils/optimize_utils.py b/src/sagemaker/serve/utils/optimize_utils.py
@@ -259,8 +259,9 @@ def _is_s3_uri(s3_uri: Optional[str]) -> bool:
 
 
 def _extract_optimization_config_and_env(
-    quantization_config: Optional[Dict] = None, compilation_config: Optional[Dict] = None,
-    sharding_config: Optional[Dict] = None
+    quantization_config: Optional[Dict] = None,
+    compilation_config: Optional[Dict] = None,
+    sharding_config: Optional[Dict] = None,
 ) -> Optional[Tuple[Optional[Dict], Optional[Dict]]]:
     """Extracts optimization config and environment variables.
 
@@ -282,9 +283,7 @@ def _extract_optimization_config_and_env(
             "OverrideEnvironment"
         )
     if sharding_config:
-        return {"ModelShardingConfig": sharding_config}, sharding_config.get(
-            "OverrideEnvironment"
-        )
+        return {"ModelShardingConfig": sharding_config}, sharding_config.get("OverrideEnvironment")
     return None, None
 
 
diff --git a/src/sagemaker/serve/validations/optimization.py b/src/sagemaker/serve/validations/optimization.py
@@ -10,155 +10,175 @@
 # 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.
-"""Holds the validation logic used for the .optimize() function"""
+"""Holds the validation logic used for the .optimize() function. INTERNAL only"""
+from __future__ import absolute_import
+
+import textwrap
+import logging
 from typing import Any, Dict, Set
 from enum import Enum
 from pydantic import BaseModel
-import textwrap
-import logging
 
 logger = logging.getLogger(__name__)
 
 
-class OptimizationContainer(Enum):
+class _OptimizationContainer(Enum):
+    """Optimization containers"""
+
     TRT = "trt"
     VLLM = "vllm"
     NEURON = "neuron"
 
 
-class OptimizationCombination(BaseModel):
-    optimization_container: OptimizationContainer = None
+class _OptimizationCombination(BaseModel):
+    """Optimization ruleset data structure for comparing input to ruleset"""
+
+    optimization_container: _OptimizationContainer = None
     compilation: bool
     speculative_decoding: bool
     sharding: bool
     quantization_technique: Set[str | None]
 
-    def validate_against(self, optimization_combination, rule_set: OptimizationContainer):
+    def validate_against(self, optimization_combination, rule_set: _OptimizationContainer):
+        """Validator for optimization containers"""
+
         if not optimization_combination.compilation == self.compilation:
-            raise ValueError("model compilation is not supported")
-        if not optimization_combination.quantization_technique.issubset(self.quantization_technique):
-            raise ValueError("model quantization is not supported")
+            raise ValueError("Compilation")
+        if not optimization_combination.quantization_technique.issubset(
+            self.quantization_technique
+        ):
+            raise ValueError(
+                f"Quantization:{optimization_combination.quantization_technique.pop()}"
+            )
         if not optimization_combination.speculative_decoding == self.speculative_decoding:
-            raise ValueError("speculative decoding is not supported")
+            raise ValueError("Speculative Decoding")
         if not optimization_combination.sharding == self.sharding:
-            raise ValueError("model sharding is not supported")
+            raise ValueError("Sharding")
 
-        if rule_set == OptimizationContainer == OptimizationContainer.TRT:
-            if optimization_combination.compilation and optimization_combination.speculative_decoding:
-                raise ValueError("model compilation and speculative decoding provided together ")
+        if rule_set == _OptimizationContainer == _OptimizationContainer.TRT:
+            if (
+                optimization_combination.compilation
+                and optimization_combination.speculative_decoding
+            ):
+                raise ValueError("Compilation and Speculative Decoding")
         else:
-            if optimization_combination.compilation and optimization_combination.quantization_technique:
-                raise ValueError("model compilation and model quantization provided together is not supported")
+            if (
+                optimization_combination.compilation
+                and optimization_combination.quantization_technique
+            ):
+                raise ValueError(
+                    f"Compilation and Quantization:{optimization_combination.quantization_technique.pop()}"
+                )
 
 
 TRT_CONFIGURATION = {
     "supported_instance_families": {"p4d", "p4de", "p5", "g5", "g6"},
-    "optimization_combination": OptimizationCombination(
-        optimization_container=OptimizationContainer.TRT,
+    "optimization_combination": _OptimizationCombination(
+        optimization_container=_OptimizationContainer.TRT,
         compilation=True,
         quantization_technique={"awq", "fp8", "smooth_quant"},
         speculative_decoding=False,
         sharding=False,
-    )
+    ),
 }
 VLLM_CONFIGURATION = {
     "supported_instance_families": {"p4d", "p4de", "p5", "g5", "g6"},
-    "optimization_combination": OptimizationCombination(
-        optimization_container=OptimizationContainer.VLLM,
+    "optimization_combination": _OptimizationCombination(
+        optimization_container=_OptimizationContainer.VLLM,
         compilation=False,
         quantization_technique={"awq", "fp8"},
         speculative_decoding=True,
-        sharding=True
-    )
+        sharding=True,
+    ),
 }
 NEURON_CONFIGURATION = {
     "supported_instance_families": {"inf2", "trn1", "trn1n"},
-    "optimization_combination": OptimizationCombination(
-        optimization_container=OptimizationContainer.NEURON,
+    "optimization_combination": _OptimizationCombination(
+        optimization_container=_OptimizationContainer.NEURON,
         compilation=True,
         quantization_technique=set(),
         speculative_decoding=False,
-        sharding=False
-    )
+        sharding=False,
+    ),
 }
 
 VALIDATION_ERROR_MSG = (
-    "The model cannot be optimized with the provided configurations on "
-    "{optimization_container} supported {instance_type} because {validation_error}."
+    "Optimizations that use {optimization_technique} "
+    "are not currently supported on {instance_type} instances"
 )
 
 
-def validate_optimization_configuration(
+def _validate_optimization_configuration(
     instance_type: str,
     quantization_config: Dict[str, Any],
     compilation_config: Dict[str, Any],
     sharding_config: Dict[str, Any],
-    speculative_decoding_config: Dict[str, Any]
+    speculative_decoding_config: Dict[str, Any],
 ):
+    """Validate .optimize() input off of standard ruleset"""
+
     split_instance_type = instance_type.split(".")
     instance_family = None
     if len(split_instance_type) == 3:  # invalid instance type will be caught below
         instance_family = split_instance_type[1]
 
     if (
-        not instance_family in TRT_CONFIGURATION["supported_instance_families"] and
-        not instance_family in VLLM_CONFIGURATION["supported_instance_families"] and
-        not instance_family in NEURON_CONFIGURATION["supported_instance_families"]
+        instance_family not in TRT_CONFIGURATION["supported_instance_families"]
+        and instance_family not in VLLM_CONFIGURATION["supported_instance_families"]
+        and instance_family not in NEURON_CONFIGURATION["supported_instance_families"]
     ):
-        invalid_instance_type_msg = f"""
-        The model cannot be optimized on {instance_type}. Please optimize on the following instance type families:
-        - For {OptimizationContainer.TRT} optimized container: {TRT_CONFIGURATION["supported_instance_families"]}
-        - For {OptimizationContainer.VLLM} optimized container: {VLLM_CONFIGURATION["supported_instance_families"]}
-        - For {OptimizationContainer.NEURON} optimized container: {NEURON_CONFIGURATION["supported_instance_families"]}
-        """
-        raise ValueError(textwrap.dedent(invalid_instance_type_msg))
-
-    optimization_combination = OptimizationCombination(
+        invalid_instance_type_msg = (
+            f"Optimizations that use {instance_type} are not currently supported"
+        )
+        raise ValueError(invalid_instance_type_msg)
+
+    optimization_combination = _OptimizationCombination(
         compilation=not compilation_config,
         speculative_decoding=not speculative_decoding_config,
         sharding=not sharding_config,
-        quantization_technique={quantization_config.get("OPTION_QUANTIZE") if quantization_config else None}
+        quantization_technique={
+            quantization_config.get("OPTION_QUANTIZE") if quantization_config else None
+        },
     )
 
     if instance_type in NEURON_CONFIGURATION["supported_instance_families"]:
         try:
             (
-                NEURON_CONFIGURATION["optimization_combination"]
-                .validate_against(optimization_combination, rule_set=OptimizationContainer.VLLM)
+                NEURON_CONFIGURATION["optimization_combination"].validate_against(
+                    optimization_combination, rule_set=_OptimizationContainer.VLLM
+                )
             )
         except ValueError as neuron_compare_error:
             raise ValueError(
                 VALIDATION_ERROR_MSG.format(
-                    optimization_container=OptimizationContainer.NEURON.value,
-                    instance_type=instance_type,
-                    validation_error=neuron_compare_error
+                    optimization_container=str(neuron_compare_error),
+                    instance_type="Neuron",
                 )
             )
     else:
         try:
             (
-                TRT_CONFIGURATION["optimization_combination"]
-                .validate_against(optimization_combination, rule_set=OptimizationContainer.TRT)
+                TRT_CONFIGURATION["optimization_combination"].validate_against(
+                    optimization_combination, rule_set=_OptimizationContainer.TRT
+                )
             )
         except ValueError as trt_compare_error:
             try:
                 (
-                    VLLM_CONFIGURATION["optimization_combination"]
-                    .validate_against(optimization_combination, rule_set=OptimizationContainer.VLLM)
+                    VLLM_CONFIGURATION["optimization_combination"].validate_against(
+                        optimization_combination, rule_set=_OptimizationContainer.VLLM
+                    )
                 )
             except ValueError as vllm_compare_error:
                 trt_error_msg = VALIDATION_ERROR_MSG.format(
-                    optimization_container=OptimizationContainer.TRT.value,
-                    instance_type=instance_type,
-                    validation_error=trt_compare_error
+                    optimization_container=str(trt_compare_error), instance_type="GPU"
                 )
                 vllm_error_msg = VALIDATION_ERROR_MSG.format(
-                    optimization_container=OptimizationContainer.VLLM.value,
-                    instance_type=instance_type,
-                    validation_error=vllm_compare_error
+                    optimization_container=str(vllm_compare_error),
+                    instance_type="GPU",
                 )
                 joint_error_msg = f"""
-                The model cannot be optimized for the following reasons:
+                Optimization cannot be performed for the following reasons:
                 - {trt_error_msg}
                 - {vllm_error_msg}
                 """
diff --git a/tests/unit/sagemaker/model/test_model.py b/tests/unit/sagemaker/model/test_model.py
@@ -958,6 +958,7 @@ def test_all_framework_models_inference_component_based_endpoint_deploy_path(
         sagemaker_session.endpoint_in_service_or_not.reset_mock()
         sagemaker_session.create_model.reset_mock()
 
+
 @patch("sagemaker.utils.repack_model")
 @patch("sagemaker.fw_utils.tar_and_upload_dir")
 def test_sharded_model_force_inference_component_based_endpoint_deploy_path(
@@ -967,7 +968,7 @@ def test_sharded_model_force_inference_component_based_endpoint_deploy_path(
         HuggingFaceModel: {
             "pytorch_version": "1.7.1",
             "py_version": "py36",
-            "transformers_version": "4.6.1"
+            "transformers_version": "4.6.1",
         },
     }
 
@@ -1007,6 +1008,7 @@ def test_sharded_model_force_inference_component_based_endpoint_deploy_path(
         sagemaker_session.endpoint_in_service_or_not.reset_mock()
         sagemaker_session.create_model.reset_mock()
 
+
 @patch("sagemaker.utils.repack_model")
 def test_repack_code_location_with_key_prefix(repack_model, sagemaker_session):
 
diff --git a/tests/unit/sagemaker/serve/utils/test_optimize_utils.py b/tests/unit/sagemaker/serve/utils/test_optimize_utils.py
@@ -326,7 +326,9 @@ def test_is_s3_uri(s3_uri, expected):
 def test_extract_optimization_config_and_env(
     quantization_config, compilation_config, sharding_config, expected_config, expected_env
 ):
-    assert _extract_optimization_config_and_env(quantization_config, compilation_config, sharding_config) == (
+    assert _extract_optimization_config_and_env(
+        quantization_config, compilation_config, sharding_config
+    ) == (
         expected_config,
         expected_env,
     )
