fixing validation bugs

Author: gwang111

src/sagemaker/serve/validations/optimization.py

@@ -25,85 +25,107 @@
 class _OptimizationContainer(Enum):
     """Optimization containers"""
 
-    TRT = "trt"
-    VLLM = "vllm"
-    NEURON = "neuron"
+    TRT = "TRT"
+    VLLM = "vLLM"
+    NEURON = "Neuron"
 
 
 class _OptimizationCombination(BaseModel):
     """Optimization ruleset data structure for comparing input to ruleset"""
 
     optimization_container: _OptimizationContainer = None
-    compilation: bool
-    speculative_decoding: bool
-    sharding: bool
+    compilation: Set[bool | None]
+    speculative_decoding: Set[bool | None]
+    sharding: Set[bool | None]
     quantization_technique: Set[str | None]
 
     def validate_against(self, optimization_combination, rule_set: _OptimizationContainer):
         """Validator for optimization containers"""
 
-        if not optimization_combination.compilation == self.compilation:
+        # check the case where no optimization combination is provided
+        if (
+            optimization_combination.compilation == {None}
+            and optimization_combination.quantization_technique == {None}
+            and optimization_combination.speculative_decoding == {None}
+            and optimization_combination.sharding == {None}
+        ):
+            raise ValueError("Optimizations are not currently supported without optimization configurations.")
+
+        # check the validity of each individual field
+        if not optimization_combination.compilation.issubset(self.compilation):
             raise ValueError("Compilation")
         if not optimization_combination.quantization_technique.issubset(
             self.quantization_technique
         ):
+            copy_quantization_technique = optimization_combination.quantization_technique.copy()
             raise ValueError(
-                f"Quantization:{optimization_combination.quantization_technique.pop()}"
+                f"Quantization:{copy_quantization_technique.pop()}"
             )
-        if not optimization_combination.speculative_decoding == self.speculative_decoding:
+        if not optimization_combination.speculative_decoding.issubset(self.speculative_decoding):
             raise ValueError("Speculative Decoding")
-        if not optimization_combination.sharding == self.sharding:
+        if not optimization_combination.sharding.issubset(self.sharding):
             raise ValueError("Sharding")
 
-        if rule_set == _OptimizationContainer == _OptimizationContainer.TRT:
+        # optimization technique combinations that need to be validated
+        if rule_set == _OptimizationContainer.TRT:
             if (
                 optimization_combination.compilation
                 and optimization_combination.speculative_decoding
             ):
-                raise ValueError("Compilation and Speculative Decoding")
+                copy_compilation = optimization_combination.compilation.copy()
+                copy_speculative_decoding = optimization_combination.speculative_decoding.copy()
+                if (
+                    copy_compilation.pop()
+                    and copy_speculative_decoding.pop()
+                ):  # Check that the 2 techniques are not None
+                    raise ValueError("Compilation and Speculative Decoding")
         else:
+            copy_compilation = optimization_combination.compilation.copy()
+            copy_quantization_technique = optimization_combination.quantization_technique.copy()
             if (
-                optimization_combination.compilation
-                and optimization_combination.quantization_technique
-            ):
+                copy_compilation.pop()
+                and copy_quantization_technique.pop()
+            ):  # Check that the 2 techniques are not None
                 raise ValueError(
                     f"Compilation and Quantization:{optimization_combination.quantization_technique.pop()}"
                 )
 
 
+TRUTHY_SET = {None, True}
+FALSY_SET = {None, False}
 TRT_CONFIGURATION = {
     "supported_instance_families": {"p4d", "p4de", "p5", "g5", "g6"},
     "optimization_combination": _OptimizationCombination(
         optimization_container=_OptimizationContainer.TRT,
-        compilation=True,
-        quantization_technique={"awq", "fp8", "smooth_quant"},
-        speculative_decoding=False,
-        sharding=False,
+        compilation=TRUTHY_SET,
+        quantization_technique={None, "awq", "fp8", "smooth_quant"},
+        speculative_decoding=FALSY_SET,
+        sharding=FALSY_SET,
     ),
 }
 VLLM_CONFIGURATION = {
     "supported_instance_families": {"p4d", "p4de", "p5", "g5", "g6"},
     "optimization_combination": _OptimizationCombination(
         optimization_container=_OptimizationContainer.VLLM,
-        compilation=False,
-        quantization_technique={"awq", "fp8"},
-        speculative_decoding=True,
-        sharding=True,
+        compilation=FALSY_SET,
+        quantization_technique={None, "awq", "fp8"},
+        speculative_decoding=TRUTHY_SET,
+        sharding=TRUTHY_SET,
     ),
 }
 NEURON_CONFIGURATION = {
     "supported_instance_families": {"inf2", "trn1", "trn1n"},
     "optimization_combination": _OptimizationCombination(
         optimization_container=_OptimizationContainer.NEURON,
-        compilation=True,
-        quantization_technique=set(),
-        speculative_decoding=False,
-        sharding=False,
+        compilation=TRUTHY_SET,
+        quantization_technique={None},
+        speculative_decoding=FALSY_SET,
+        sharding=FALSY_SET,
     ),
 }
 
 VALIDATION_ERROR_MSG = (
-    "Optimizations that use {optimization_technique} "
+    "Optimizations for {optimization_container} that use {optimization_technique} "
     "are not currently supported on {instance_type} instances"
 )
 
@@ -117,28 +139,41 @@ def _validate_optimization_configuration(
 ):
     """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 instance_type:
+        split_instance_type = instance_type.split(".")
+        if len(split_instance_type) == 3:
+            instance_family = split_instance_type[1]
 
     if (
         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"Optimizations that use {instance_type} are not currently supported"
+            f"Optimizations that uses {instance_type} instance type are not currently supported"
         )
         raise ValueError(invalid_instance_type_msg)
 
+    quantization_technique = None
+    if (
+        quantization_config
+        and quantization_config.get("OverrideEnvironment")
+        and quantization_config.get("OverrideEnvironment").get("OPTION_QUANTIZE")
+    ):
+        quantization_technique = quantization_config.get("OverrideEnvironment").get("OPTION_QUANTIZE")
+
     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
+        compilation={
+            None if compilation_config is None else bool(compilation_config)
+        },
+        speculative_decoding={
+            None if speculative_decoding_config is None else bool(speculative_decoding_config)
+        },
+        sharding={
+            None if sharding_config is None else bool(sharding_config)
         },
+        quantization_technique={quantization_technique},
     )
 
     if instance_type in NEURON_CONFIGURATION["supported_instance_families"]:
@@ -151,7 +186,8 @@ def _validate_optimization_configuration(
         except ValueError as neuron_compare_error:
             raise ValueError(
                 VALIDATION_ERROR_MSG.format(
-                    optimization_container=str(neuron_compare_error),
+                    optimization_container=_OptimizationContainer.NEURON.value,
+                    optimization_technique=str(neuron_compare_error),
                     instance_type="Neuron",
                 )
             )
@@ -171,10 +207,13 @@ def _validate_optimization_configuration(
                 )
             except ValueError as vllm_compare_error:
                 trt_error_msg = VALIDATION_ERROR_MSG.format(
-                    optimization_container=str(trt_compare_error), instance_type="GPU"
+                    optimization_container=_OptimizationContainer.TRT.value,
+                    optimization_technique=str(trt_compare_error),
+                    instance_type="GPU"
                 )
                 vllm_error_msg = VALIDATION_ERROR_MSG.format(
-                    optimization_container=str(vllm_compare_error),
+                    optimization_container=_OptimizationContainer.VLLM.value,
+                    optimization_technique=str(vllm_compare_error),
                     instance_type="GPU",
                 )
                 joint_error_msg = f"""

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

@@ -11,6 +11,8 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
+
+import textwrap
 from unittest.mock import MagicMock, patch, Mock, mock_open
 
 import unittest
@@ -25,6 +27,7 @@
 from sagemaker.serve.utils.exceptions import TaskNotFoundException
 from sagemaker.serve.utils.predictors import TensorflowServingLocalPredictor
 from sagemaker.serve.utils.types import ModelServer
+from sagemaker.serve.validations.optimization import _validate_optimization_configuration
 from tests.unit.sagemaker.serve.constants import MOCK_IMAGE_CONFIG, MOCK_VPC_CONFIG
 
 schema_builder = MagicMock()
@@ -2383,7 +2386,7 @@ def test_optimize(
         builder.pysdk_model = pysdk_model
 
         job_name = "my-optimization-job"
-        instance_type = "ml.inf1.xlarge"
+        instance_type = "ml.inf2.xlarge"
         output_path = "s3://my-bucket/output"
         quantization_config = {
             "Image": "quantization-image-uri",
@@ -2425,7 +2428,7 @@ def test_optimize(
         mock_send_telemetry.assert_called_once()
         mock_sagemaker_session.sagemaker_client.create_optimization_job.assert_called_once_with(
             OptimizationJobName="my-optimization-job",
-            DeploymentInstanceType="ml.inf1.xlarge",
+            DeploymentInstanceType="ml.inf2.xlarge",
             RoleArn="arn:aws:iam::123456789012:role/SageMakerRole",
             OptimizationEnvironment={"Var1": "value", "Var2": "value"},
             ModelSource={"S3": {"S3Uri": "s3://uri"}},
@@ -2646,6 +2649,7 @@ def test_optimize_local_mode(self, mock_get_serve_setting):
             ValueError,
             "Model optimization is only supported in Sagemaker Endpoint Mode.",
             lambda: model_builder.optimize(
+                instance_type="ml.g5.24xlarge",
                 quantization_config={"OverrideEnvironment": {"OPTION_QUANTIZE": "awq"}}
             ),
         )
@@ -2662,6 +2666,7 @@ def test_optimize_exclusive_args(self, mock_get_serve_setting):
             ValueError,
             "Quantization config and compilation config are mutually exclusive.",
             lambda: model_builder.optimize(
+                instance_type="ml.g5.24xlarge",
                 quantization_config={"OverrideEnvironment": {"OPTION_QUANTIZE": "awq"}},
                 compilation_config={"OverrideEnvironment": {"OPTION_QUANTIZE": "awq"}},
             ),
@@ -2675,10 +2680,17 @@ def test_optimize_exclusive_sharding(self, mock_get_serve_setting):
             sagemaker_session=mock_sagemaker_session,
         )
 
+        expected_error_message = """
+        Optimization cannot be performed for the following reasons:
+        - Optimizations for TRT that use Sharding are not currently supported on GPU instances
+        - Optimizations for vLLM that use Compilation are not currently supported on GPU instances
+        """
+
         self.assertRaisesRegex(
             ValueError,
-            "Sharding config is mutually exclusive and cannot be combined with any other optimization.",
+            textwrap.dedent(expected_error_message),
             lambda: model_builder.optimize(
+                instance_type="ml.g5.24xlarge",
                 quantization_config={"OverrideEnvironment": {"OPTION_QUANTIZE": "awq"}},
                 compilation_config={"OverrideEnvironment": {"OPTION_QUANTIZE": "awq"}},
                 sharding_config={"OverrideEnvironment": {"OPTION_QUANTIZE": "awq"}},
@@ -2697,6 +2709,7 @@ def test_optimize_exclusive_sharding_args(self, mock_get_serve_setting):
             ValueError,
             "OPTION_TENSOR_PARALLEL_DEGREE is required environment variable with Sharding config.",
             lambda: model_builder.optimize(
+                instance_type="ml.g5.24xlarge",
                 sharding_config={"OverrideEnvironment": {"OPTION_QUANTIZE": "awq"}},
             ),
         )
@@ -2820,3 +2833,69 @@ def test_optimize_for_hf_without_custom_s3_path(
                 "OutputConfig": {"S3OutputLocation": "s3://bucket/code/"},
             },
         )
+
+
+class TestModelBuilderOptimizeValidations(unittest.TestCase):
+
+    def test_corner_cases_throw_errors(self):
+        self.assertRaisesRegex(
+            ValueError,
+            "Optimizations that uses None instance type are not currently supported",
+            lambda: _validate_optimization_configuration(
+                sharding_config={"OverrideEnvironment": {"OPTION_QUANTIZE": "awq"}},
+                instance_type=None,
+                quantization_config=None,
+                speculative_decoding_config=None,
+                compilation_config=None,
+            ),
+        )
+        self.assertRaisesRegex(
+            ValueError,
+            "Optimizations are not currently supported without optimization configurations.",
+            lambda: _validate_optimization_configuration(
+                instance_type="ml.g5.24xlarge",
+                quantization_config=None,
+                speculative_decoding_config=None,
+                compilation_config=None,
+                sharding_config=None,
+            ),
+        )
+
+    def test_trt_and_vllm_configurations_throw_errors_for_rule_set(self):
+        expected_quantization_error_message = """
+        Optimization cannot be performed for the following reasons:
+        - Optimizations for TRT that use Quantization:test are not currently supported on GPU instances
+        - Optimizations for vLLM that use Quantization:test are not currently supported on GPU instances
+        """
+        self.assertRaisesRegex(
+            ValueError,
+            textwrap.dedent(expected_quantization_error_message),
+            lambda: _validate_optimization_configuration(
+                instance_type="ml.g5.24xlarge",
+                quantization_config={
+                    "OverrideEnvironment": {"OPTION_QUANTIZE": "test"},
+                },
+                sharding_config=None,
+                speculative_decoding_config=None,
+                compilation_config=None,
+            ),
+        )
+
+    @patch.object(ModelBuilder, "_get_serve_setting", autospec=True)
+    def test_neuron_configurations_throw_errors_for_rule_set(self, mock_get_serve_setting):
+        pass
+
+    def test_trt_configurations_rule_set(self):
+        _validate_optimization_configuration(
+            instance_type="ml.g5.24xlarge",
+            quantization_config={
+                "OverrideEnvironment": {"OPTION_QUANTIZE": "awq"},
+            },
+            sharding_config=None,
+            speculative_decoding_config=None,
+            compilation_config=None,
+        )
+
+    @patch.object(ModelBuilder, "_get_serve_setting", autospec=True)
+    def test_vllm_configurations_rule_set(self, mock_get_serve_setting):
+        pass