add UTs

Author: gwang111

src/sagemaker/serve/validations/optimization.py

@@ -41,6 +41,9 @@ class _OptimizationCombination(BaseModel):
 
     def validate_against(self, optimization_combination, rule_set: _OptimizationContainer):
         """Validator for optimization containers"""
+        print(optimization_combination)
+        print(rule_set)
+        print(optimization_combination.speculative_decoding.issubset(self.speculative_decoding))
 
         # check the case where no optimization combination is provided
         if (
@@ -49,7 +52,7 @@ def validate_against(self, optimization_combination, rule_set: _OptimizationCont
             and optimization_combination.speculative_decoding == {None}
             and optimization_combination.sharding == {None}
         ):
-            raise ValueError("Optimizations are not currently supported without optimization configurations.")
+            raise ValueError("no optimization configurations")
 
         # check the validity of each individual field
         if not optimization_combination.compilation.issubset(self.compilation):
@@ -58,9 +61,7 @@ def validate_against(self, optimization_combination, rule_set: _OptimizationCont
             self.quantization_technique
         ):
             copy_quantization_technique = optimization_combination.quantization_technique.copy()
-            raise ValueError(
-                f"Quantization:{copy_quantization_technique.pop()}"
-            )
+            raise ValueError(f"Quantization:{copy_quantization_technique.pop()}")
         if not optimization_combination.speculative_decoding.issubset(self.speculative_decoding):
             raise ValueError("Speculative Decoding")
         if not optimization_combination.sharding.issubset(self.sharding):
@@ -75,16 +76,14 @@ def validate_against(self, optimization_combination, rule_set: _OptimizationCont
                 copy_compilation = optimization_combination.compilation.copy()
                 copy_speculative_decoding = optimization_combination.speculative_decoding.copy()
                 if (
-                    copy_compilation.pop()
-                    and copy_speculative_decoding.pop()
+                    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 (
-                copy_compilation.pop()
-                and copy_quantization_technique.pop()
+                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()}"
@@ -161,26 +160,24 @@ def _validate_optimization_configuration(
         and quantization_config.get("OverrideEnvironment")
         and quantization_config.get("OverrideEnvironment").get("OPTION_QUANTIZE")
     ):
-        quantization_technique = quantization_config.get("OverrideEnvironment").get("OPTION_QUANTIZE")
+        quantization_technique = quantization_config.get("OverrideEnvironment").get(
+            "OPTION_QUANTIZE"
+        )
 
     optimization_combination = _OptimizationCombination(
-        compilation={
-            None if compilation_config is None else bool(compilation_config)
-        },
+        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)
-        },
+        sharding={None if sharding_config is None else bool(sharding_config)},
         quantization_technique={quantization_technique},
     )
 
-    if instance_type in NEURON_CONFIGURATION["supported_instance_families"]:
+    if instance_family in NEURON_CONFIGURATION["supported_instance_families"]:
         try:
             (
                 NEURON_CONFIGURATION["optimization_combination"].validate_against(
-                    optimization_combination, rule_set=_OptimizationContainer.VLLM
+                    optimization_combination, rule_set=_OptimizationContainer.NEURON
                 )
             )
         except ValueError as neuron_compare_error:
@@ -209,7 +206,7 @@ def _validate_optimization_configuration(
                 trt_error_msg = VALIDATION_ERROR_MSG.format(
                     optimization_container=_OptimizationContainer.TRT.value,
                     optimization_technique=str(trt_compare_error),
-                    instance_type="GPU"
+                    instance_type="GPU",
                 )
                 vllm_error_msg = VALIDATION_ERROR_MSG.format(
                     optimization_container=_OptimizationContainer.VLLM.value,

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

@@ -2386,11 +2386,11 @@ def test_optimize(
         builder.pysdk_model = pysdk_model
 
         job_name = "my-optimization-job"
-        instance_type = "ml.inf2.xlarge"
+        instance_type = "ml.g5.24xlarge"
         output_path = "s3://my-bucket/output"
         quantization_config = {
             "Image": "quantization-image-uri",
-            "OverrideEnvironment": {"ENV_VAR": "value"},
+            "OverrideEnvironment": {"OPTION_QUANTIZE": "awq"},
         }
         env_vars = {"Var1": "value", "Var2": "value"}
         kms_key = "arn:aws:kms:us-west-2:123456789012:key/my-key-id"
@@ -2428,15 +2428,15 @@ 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.inf2.xlarge",
+            DeploymentInstanceType="ml.g5.24xlarge",
             RoleArn="arn:aws:iam::123456789012:role/SageMakerRole",
             OptimizationEnvironment={"Var1": "value", "Var2": "value"},
             ModelSource={"S3": {"S3Uri": "s3://uri"}},
             OptimizationConfigs=[
                 {
                     "ModelQuantizationConfig": {
                         "Image": "quantization-image-uri",
-                        "OverrideEnvironment": {"ENV_VAR": "value"},
+                        "OverrideEnvironment": {"OPTION_QUANTIZE": "awq"},
                     }
                 }
             ],
@@ -2650,7 +2650,7 @@ def test_optimize_local_mode(self, mock_get_serve_setting):
             "Model optimization is only supported in Sagemaker Endpoint Mode.",
             lambda: model_builder.optimize(
                 instance_type="ml.g5.24xlarge",
-                quantization_config={"OverrideEnvironment": {"OPTION_QUANTIZE": "awq"}}
+                quantization_config={"OverrideEnvironment": {"OPTION_QUANTIZE": "awq"}},
             ),
         )
 
@@ -2842,16 +2842,22 @@ def test_corner_cases_throw_errors(self):
             ValueError,
             "Optimizations that uses None instance type are not currently supported",
             lambda: _validate_optimization_configuration(
-                sharding_config={"OverrideEnvironment": {"OPTION_QUANTIZE": "awq"}},
+                sharding_config={"key": "value"},
                 instance_type=None,
                 quantization_config=None,
                 speculative_decoding_config=None,
                 compilation_config=None,
             ),
         )
+
+        expected_missing_optimization_configs_error_message = """
+        Optimization cannot be performed for the following reasons:
+        - Optimizations for TRT that use no optimization configurations are not currently supported on GPU instances
+        - Optimizations for vLLM that use no optimization configurations are not currently supported on GPU instances
+        """
         self.assertRaisesRegex(
             ValueError,
-            "Optimizations are not currently supported without optimization configurations.",
+            textwrap.dedent(expected_missing_optimization_configs_error_message),
             lambda: _validate_optimization_configuration(
                 instance_type="ml.g5.24xlarge",
                 quantization_config=None,
@@ -2881,11 +2887,39 @@ def test_trt_and_vllm_configurations_throw_errors_for_rule_set(self):
             ),
         )
 
-    @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_neuron_configurations_throw_errors_for_rule_set(self):
+        self.assertRaisesRegex(
+            ValueError,
+            (
+                "Optimizations for Neuron that use Speculative Decoding "
+                "are not currently supported on Neuron instances"
+            ),
+            lambda: _validate_optimization_configuration(
+                instance_type="ml.inf2.xlarge",
+                quantization_config=None,
+                speculative_decoding_config={"key": "value"},
+                compilation_config=None,
+                sharding_config=None,
+            ),
+        )
+
+        self.assertRaisesRegex(
+            ValueError,
+            (
+                "Optimizations for Neuron that use Sharding "
+                "are not currently supported on Neuron instances"
+            ),
+            lambda: _validate_optimization_configuration(
+                instance_type="ml.inf2.xlarge",
+                quantization_config=None,
+                speculative_decoding_config=None,
+                compilation_config=None,
+                sharding_config={"key": "value"},
+            ),
+        )
 
     def test_trt_configurations_rule_set(self):
+        # Can be quantized
         _validate_optimization_configuration(
             instance_type="ml.g5.24xlarge",
             quantization_config={
@@ -2896,6 +2930,51 @@ def test_trt_configurations_rule_set(self):
             compilation_config=None,
         )
 
-    @patch.object(ModelBuilder, "_get_serve_setting", autospec=True)
-    def test_vllm_configurations_rule_set(self, mock_get_serve_setting):
-        pass
+        # 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"},
+        )
+
+    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",
+            quantization_config=None,
+            sharding_config=None,
+            speculative_decoding_config={"key": "value"},
+            compilation_config=None,
+        )
+
+        # Can be sharded
+        _validate_optimization_configuration(
+            instance_type="ml.g5.24xlarge",
+            quantization_config=None,
+            sharding_config={"key": "value"},
+            speculative_decoding_config=None,
+            compilation_config=None,
+        )
+
+    def test_neuron_configurations_rule_set(self):
+        # Can be compiled
+        _validate_optimization_configuration(
+            instance_type="ml.inf2.xlarge",
+            quantization_config=None,
+            sharding_config=None,
+            speculative_decoding_config=None,
+            compilation_config={"key": "value"},
+        )