fix UTs

Author: gwang111

src/sagemaker/serve/validations/optimization.py

@@ -57,26 +57,21 @@ def validate_against(self, optimization_combination, rule_set: _OptimizationCont
 
         # optimization technique combinations that need to be validated
         if optimization_combination.compilation and optimization_combination.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
+            is_compiled = optimization_combination.compilation.copy().pop()
+            is_speculative_decoding = optimization_combination.speculative_decoding.copy().pop()
+            if is_compiled and is_speculative_decoding:
                 raise ValueError("Compilation and Speculative Decoding together")
 
         if rule_set == _OptimizationContainer.TRT:
-            if (
-                optimization_combination.compilation
-                and not optimization_combination.quantization_technique
-                or not optimization_combination.compilation
-                and optimization_combination.quantization_technique
-            ):
+            is_compiled = optimization_combination.compilation.copy().pop()
+            is_quantized = optimization_combination.quantization_technique.copy().pop()
+            if is_compiled and not is_quantized or is_quantized and not is_compiled:
                 raise ValueError("Compilation must be provided with Quantization")
         else:
-            copy_compilation = optimization_combination.compilation.copy()
-            copy_quantization_technique = optimization_combination.quantization_technique.copy()
+            is_compiled = optimization_combination.compilation.copy().pop()
+            is_quantization_technique = optimization_combination.quantization_technique.copy().pop()
             if (
-                copy_compilation.pop() and copy_quantization_technique.pop()
+                is_compiled and is_quantization_technique
             ):  # Check that the 2 techniques are not None
                 raise ValueError(
                     f"Compilation and Quantization:{optimization_combination.quantization_technique.pop()}"
@@ -99,8 +94,8 @@ def validate_against(self, optimization_combination, rule_set: _OptimizationCont
     "supported_instance_families": {"p4d", "p4de", "p5", "g5", "g6"},
     "optimization_combination": _OptimizationCombination(
         optimization_container=_OptimizationContainer.VLLM,
-        compilation=TRUTHY_SET,
-        quantization_technique={None, "awq", "fp8"},
+        compilation=FALSY_SET,
+        quantization_technique={None},
         speculative_decoding=TRUTHY_SET,
         sharding=TRUTHY_SET,
     ),
@@ -203,8 +198,9 @@ def _validate_optimization_configuration(
                         optimization_combination, rule_set=_OptimizationContainer.VLLM
                     )
                 )
+                print("fsdafas")
             except ValueError as vllm_compare_error:
-                if trt_compare_error == "Compilation must be provided with Quantization":
+                if str(trt_compare_error) == "Compilation must be provided with Quantization":
                     joint_error_msg = f"""
                     Optimization cannot be performed for the following reasons:
                     - Optimizations that use {trt_compare_error} and vice-versa for GPU instances.

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

@@ -2912,37 +2912,44 @@ def test_neuron_configurations_throw_errors_for_rule_set(self):
 
     def test_trt_configurations_rule_set(self):
         # Can be quantized
-        _validate_optimization_configuration(
-            instance_type="ml.g5.24xlarge",
-            quantization_config={
-                "OverrideEnvironment": {"OPTION_QUANTIZE": "awq"},
-            },
-            sharding_config=None,
-            speculative_decoding_config=None,
-            compilation_config=None,
+        expected_compilation_quantization_error_message = """
+        Optimization cannot be performed for the following reasons:
+        - Optimizations that use Compilation must be provided with Quantization and vice-versa for GPU instances.
+        - Optimizations that use Quantization:awq are not supported for GPU instances.
+        """
+        self.assertRaisesRegex(
+            ValueError,
+            textwrap.dedent(expected_compilation_quantization_error_message),
+            lambda: _validate_optimization_configuration(
+                instance_type="ml.g5.24xlarge",
+                quantization_config={
+                    "OverrideEnvironment": {"OPTION_QUANTIZE": "awq"},
+                },
+                sharding_config=None,
+                speculative_decoding_config=None,
+                compilation_config=None,
+            ),
         )
 
         # Can be compiled
-        _validate_optimization_configuration(
-            instance_type="ml.g5.24xlarge",
-            quantization_config=None,
-            sharding_config=None,
-            speculative_decoding_config=None,
-            compilation_config={"key": "value"},
+        expected_compilation_quantization_error_message = """
+        Optimization cannot be performed for the following reasons:
+        - Optimizations that use Compilation must be provided with Quantization and vice-versa for GPU instances.
+        - Optimizations that use Compilation are not supported for GPU instances.
+        """
+        self.assertRaisesRegex(
+            ValueError,
+            textwrap.dedent(expected_compilation_quantization_error_message),
+            lambda: _validate_optimization_configuration(
+                instance_type="ml.g5.24xlarge",
+                quantization_config=None,
+                sharding_config=None,
+                speculative_decoding_config=None,
+                compilation_config={"key": "value"},
+            ),
         )
 
     def test_vllm_configurations_rule_set(self):
-        # Can be quantized
-        _validate_optimization_configuration(
-            instance_type="ml.g5.24xlarge",
-            quantization_config={
-                "OverrideEnvironment": {"OPTION_QUANTIZE": "awq"},
-            },
-            sharding_config=None,
-            speculative_decoding_config=None,
-            compilation_config=None,
-        )
-
         # Can use speculative decoding
         _validate_optimization_configuration(
             instance_type="ml.g5.24xlarge",