fix rebase issues

Author: gwang111

src/sagemaker/serve/builder/model_builder.py

@@ -1248,6 +1248,7 @@ def _model_builder_optimize_wrapper(
         # TODO: ideally these dictionaries need to be sagemaker_core shapes
         # TODO: for organization, abstract all validation behind this fn
         _validate_optimization_configuration(
+            is_jumpstart=self._is_jumpstart_model_id(),
             instance_type=instance_type,
             quantization_config=quantization_config,
             compilation_config=compilation_config,
@@ -1264,13 +1265,6 @@ def _model_builder_optimize_wrapper(
         if self.mode != Mode.SAGEMAKER_ENDPOINT:
             raise ValueError("Model optimization is only supported in Sagemaker Endpoint Mode.")
 
-        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
         ):
@@ -1456,7 +1450,9 @@ def _optimize_for_hf(
                 quantization_override_env,
                 compilation_override_env,
                 sharding_override_env,
-            ) = _extract_optimization_config_and_env(quantization_config, compilation_config)
+            ) = _extract_optimization_config_and_env(
+                quantization_config, compilation_config, sharding_config
+            )
             create_optimization_job_args["OptimizationConfigs"] = [
                 {k: v} for k, v in optimization_config.items()
             ]

src/sagemaker/serve/utils/optimize_utils.py

@@ -405,6 +405,7 @@ def _extract_optimization_config_and_env(
 
     return None, None, None, None
 
+
 def _custom_speculative_decoding(
     model: Model,
     speculative_decoding_config: Optional[Dict],

src/sagemaker/serve/validations/optimization.py

@@ -104,6 +104,7 @@ def validate_against(self, optimization_combination, rule_set: _OptimizationCont
 
 
 def _validate_optimization_configuration(
+    is_jumpstart: bool,
     instance_type: str,
     quantization_config: Dict[str, Any],
     compilation_config: Dict[str, Any],
@@ -153,6 +154,9 @@ def _validate_optimization_configuration(
         and optimization_combination.speculative_decoding == {None}
         and optimization_combination.sharding == {None}
     ):
+        # JumpStart has defaults for Inf/Trn instances
+        if is_jumpstart and instance_family in NEURON_CONFIGURATION["supported_instance_families"]:
+            return
         raise ValueError(
             (
                 "Optimizations that provide no optimization configs "

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

@@ -2927,6 +2927,7 @@ def test_optimize_with_gpu_instance_and_llama_3_1_and_compilation(
             "Compilation is not supported for Llama-3.1 with a GPU instance.",
             lambda: model_builder.optimize(
                 job_name="job_name-123",
+                instance_type="ml.g5.24xlarge",
                 compilation_config={"OverrideEnvironment": {"OPTION_TENSOR_PARALLEL_DEGREE": "2"}},
                 output_path="s3://bucket/code/",
             ),
@@ -2975,9 +2976,10 @@ def test_optimize_with_gpu_instance_and_compilation_with_speculative_decoding(
 
         self.assertRaisesRegex(
             ValueError,
-            "Compilation is not supported with speculative decoding with a GPU instance.",
+            "Optimizations that use Compilation and Speculative Decoding are not supported for GPU instances.",
             lambda: model_builder.optimize(
                 job_name="job_name-123",
+                instance_type="ml.g5.24xlarge",
                 speculative_decoding_config={
                     "ModelProvider": "custom",
                     "ModelSource": "s3://data-source",
@@ -3481,6 +3483,7 @@ def test_corner_cases_throw_errors(self):
             ValueError,
             "Optimizations that uses None instance type are not currently supported",
             lambda: _validate_optimization_configuration(
+                is_jumpstart=False,
                 sharding_config={"key": "value"},
                 instance_type=None,
                 quantization_config=None,
@@ -3496,6 +3499,7 @@ def test_corner_cases_throw_errors(self):
                 "are currently not support on both GPU and Neuron instances."
             ),
             lambda: _validate_optimization_configuration(
+                is_jumpstart=False,
                 instance_type="ml.g5.24xlarge",
                 quantization_config=None,
                 speculative_decoding_config=None,
@@ -3504,12 +3508,22 @@ def test_corner_cases_throw_errors(self):
             ),
         )
 
+        _validate_optimization_configuration(
+            is_jumpstart=True,
+            instance_type="ml.inf2.xlarge",
+            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):
         # Quantization:smoothquant without compilation
         self.assertRaisesRegex(
             ValueError,
             "Optimizations that use Quantization:smoothquant must be provided with Compilation for GPU instances.",
             lambda: _validate_optimization_configuration(
+                is_jumpstart=False,
                 instance_type="ml.g5.24xlarge",
                 quantization_config={
                     "OverrideEnvironment": {"OPTION_QUANTIZE": "smoothquant"},
@@ -3525,6 +3539,7 @@ def test_trt_and_vllm_configurations_throw_errors_for_rule_set(self):
             ValueError,
             "Optimizations that use Quantization:test are not supported for GPU instances.",
             lambda: _validate_optimization_configuration(
+                is_jumpstart=False,
                 instance_type="ml.g5.24xlarge",
                 quantization_config={
                     "OverrideEnvironment": {"OPTION_QUANTIZE": "test"},
@@ -3540,6 +3555,7 @@ def test_neuron_configurations_throw_errors_for_rule_set(self):
             ValueError,
             "Optimizations that use Speculative Decoding are not supported on Neuron instances.",
             lambda: _validate_optimization_configuration(
+                is_jumpstart=False,
                 instance_type="ml.inf2.xlarge",
                 quantization_config=None,
                 speculative_decoding_config={"key": "value"},
@@ -3552,6 +3568,7 @@ def test_neuron_configurations_throw_errors_for_rule_set(self):
             ValueError,
             "Optimizations that use Sharding are not supported on Neuron instances.",
             lambda: _validate_optimization_configuration(
+                is_jumpstart=False,
                 instance_type="ml.inf2.xlarge",
                 quantization_config=None,
                 speculative_decoding_config=None,
@@ -3563,6 +3580,7 @@ def test_neuron_configurations_throw_errors_for_rule_set(self):
     def test_trt_configurations_rule_set(self):
         # Can be compiled with quantization
         _validate_optimization_configuration(
+            is_jumpstart=False,
             instance_type="ml.g5.24xlarge",
             quantization_config={
                 "OverrideEnvironment": {"OPTION_QUANTIZE": "smoothquant"},
@@ -3574,6 +3592,7 @@ def test_trt_configurations_rule_set(self):
 
         # Can be just compiled
         _validate_optimization_configuration(
+            is_jumpstart=False,
             instance_type="ml.g5.24xlarge",
             quantization_config=None,
             sharding_config=None,
@@ -3583,6 +3602,7 @@ def test_trt_configurations_rule_set(self):
 
         # Can be just compiled with empty dict
         _validate_optimization_configuration(
+            is_jumpstart=False,
             instance_type="ml.g5.24xlarge",
             quantization_config=None,
             sharding_config=None,
@@ -3593,6 +3613,7 @@ def test_trt_configurations_rule_set(self):
     def test_vllm_configurations_rule_set(self):
         # Can use speculative decoding
         _validate_optimization_configuration(
+            is_jumpstart=False,
             instance_type="ml.g5.24xlarge",
             quantization_config=None,
             sharding_config=None,
@@ -3602,6 +3623,7 @@ def test_vllm_configurations_rule_set(self):
 
         # Can be quantized
         _validate_optimization_configuration(
+            is_jumpstart=False,
             instance_type="ml.g5.24xlarge",
             quantization_config={
                 "OverrideEnvironment": {"OPTION_QUANTIZE": "awq"},
@@ -3613,6 +3635,7 @@ def test_vllm_configurations_rule_set(self):
 
         # Can be sharded
         _validate_optimization_configuration(
+            is_jumpstart=False,
             instance_type="ml.g5.24xlarge",
             quantization_config=None,
             sharding_config={"key": "value"},
@@ -3623,6 +3646,7 @@ def test_vllm_configurations_rule_set(self):
     def test_neuron_configurations_rule_set(self):
         # Can be compiled
         _validate_optimization_configuration(
+            is_jumpstart=False,
             instance_type="ml.inf2.xlarge",
             quantization_config=None,
             sharding_config=None,
@@ -3632,6 +3656,7 @@ def test_neuron_configurations_rule_set(self):
 
         # Can be compiled with empty dict
         _validate_optimization_configuration(
+            is_jumpstart=False,
             instance_type="ml.inf2.xlarge",
             quantization_config=None,
             sharding_config=None,

tests/unit/sagemaker/serve/utils/test_optimize_utils.py

@@ -284,7 +284,10 @@ def test_is_draft_model_gated(draft_model_config, expected):
 
 
 @pytest.mark.parametrize(
-    "quantization_config, compilation_config, sharding_config, expected_config, expected_quant_env, expected_compilation_env, expected_sharding_env",
+    (
+        "quantization_config, compilation_config, sharding_config, expected_config, "
+        "expected_quant_env, expected_compilation_env, expected_sharding_env"
+    ),
     [
         (
             None,