update validation logic

Author: gwang111

src/sagemaker/serve/validations/optimization.py

@@ -176,36 +176,49 @@ def _validate_optimization_configuration(
                 )
             )
     else:
-        try:
-            (
+        if optimization_combination.compilation.copy().pop():  # Compilation is only enabled for TRT
+            try:
                 TRT_CONFIGURATION["optimization_combination"].validate_against(
                     optimization_combination, rule_set=_OptimizationContainer.TRT
                 )
-            )
-        except ValueError as trt_compare_error:
+            except ValueError as trt_compare_error:
+                raise ValueError(
+                    (
+                        f"Optimizations that use Compilation and {trt_compare_error} "
+                        "are not supported for GPU instances."
+                    )
+                )
+        else:
             try:
                 (
                     VLLM_CONFIGURATION["optimization_combination"].validate_against(
                         optimization_combination, rule_set=_OptimizationContainer.VLLM
                     )
                 )
             except ValueError as vllm_compare_error:
-                if "Quantization must be provided with Compilation" in str(trt_compare_error):
-                    joint_error_msg = f"""
+                try:  # try both VLLM and TRT to cover both rule sets
+                    (
+                        TRT_CONFIGURATION["optimization_combination"].validate_against(
+                            optimization_combination, rule_set=_OptimizationContainer.TRT
+                        )
+                    )
+                except ValueError as trt_compare_error:
+                    if "Quantization must be provided with Compilation" in str(trt_compare_error):
+                        joint_error_msg = f"""
                     Optimization cannot be performed for the following reasons:
                     - Optimizations that use {trt_compare_error} for GPU instances.
                     - Optimizations that use {vllm_compare_error} are not supported for GPU instances.
                     """
-                else:
-                    if str(trt_compare_error) == str(vllm_compare_error):
-                        joint_error_msg = f"""
+                    else:
+                        if str(trt_compare_error) == str(vllm_compare_error):
+                            joint_error_msg = f"""
                         Optimization cannot be performed for the following reasons:
                         - Optimizations that use {trt_compare_error} are not supported for GPU instances.
                         """
-                    else:
-                        joint_error_msg = f"""
+                        else:
+                            joint_error_msg = f"""
                         Optimization cannot be performed for the following reasons:
                         - Optimizations that use {trt_compare_error} are not supported for GPU instances.
                         - Optimizations that use {vllm_compare_error} are not supported for GPU instances.
                         """
-                raise ValueError(textwrap.dedent(joint_error_msg))
+                    raise ValueError(textwrap.dedent(joint_error_msg))

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

@@ -2680,15 +2680,9 @@ 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 that use Sharding are not supported for GPU instances.
-        - Optimizations that use Compilation are not supported for GPU instances.
-        """
-
         self.assertRaisesRegex(
             ValueError,
-            textwrap.dedent(expected_error_message),
+            "Optimizations that use Compilation and Sharding are not supported for GPU instances.",
             lambda: model_builder.optimize(
                 instance_type="ml.g5.24xlarge",
                 quantization_config={"OverrideEnvironment": {"OPTION_QUANTIZE": "awq"}},