Fix quantization/compilation config handling for optimize().

Author: Joseph Zhang

src/sagemaker/serve/builder/jumpstart_builder.py

@@ -727,8 +727,8 @@ def _optimize_for_jumpstart(
             pysdk_model_env_vars = self._get_neuron_model_env_vars(instance_type)
 
         # optimization_config can contain configs for both quantization and compilation
-        optimization_config, override_env = _extract_optimization_config_and_env(
-            quantization_config, compilation_config
+        optimization_config, quantization_override_env, compilation_override_env = (
+            _extract_optimization_config_and_env(quantization_config, compilation_config)
         )
         if (
             not optimization_config or not optimization_config.get("ModelCompilationConfig")
@@ -738,13 +738,16 @@ def _optimize_for_jumpstart(
                 optimization_config = {}
 
             # Fallback to default if override_env is None or empty
-            if not override_env:
-                override_env = pysdk_model_env_vars
+            if not compilation_override_env:
+                compilation_override_env = pysdk_model_env_vars
 
             # Update optimization_config with ModelCompilationConfig
-            optimization_config["ModelCompilationConfig"] = {
-                "OverrideEnvironment": override_env,
-            }
+            override_compilation_config = (
+                {"OverrideEnvironment": compilation_override_env}
+                if compilation_override_env
+                else {}
+            )
+            optimization_config["ModelCompilationConfig"] = override_compilation_config
 
         if speculative_decoding_config:
             self._set_additional_model_source(speculative_decoding_config)
@@ -798,7 +801,13 @@ def _optimize_for_jumpstart(
                     "AcceptEula": True
                 }
 
-        optimization_env_vars = _update_environment_variables(optimization_env_vars, override_env)
+        optimization_env_vars = _update_environment_variables(
+            optimization_env_vars,
+            {
+                **(quantization_override_env or {}),
+                **(compilation_override_env or {}),
+            },
+        )
         if optimization_env_vars:
             self.pysdk_model.env.update(optimization_env_vars)
         if quantization_config or is_compilation:

src/sagemaker/serve/builder/model_builder.py

@@ -1361,13 +1361,18 @@ def _optimize_for_hf(
             model_source = _generate_model_source(self.pysdk_model.model_data, False)
             create_optimization_job_args["ModelSource"] = model_source
 
-            optimization_config, override_env = _extract_optimization_config_and_env(
-                quantization_config, compilation_config
+            optimization_config, quantization_override_env, compilation_override_env = (
+                _extract_optimization_config_and_env(quantization_config, compilation_config)
             )
             create_optimization_job_args["OptimizationConfigs"] = [
                 {k: v} for k, v in optimization_config.items()
             ]
-            self.pysdk_model.env.update(override_env)
+            self.pysdk_model.env.update(
+                {
+                    **(quantization_override_env or {}),
+                    **(compilation_override_env or {}),
+                }
+            )
 
             output_config = {"S3OutputLocation": output_path}
             if kms_key:

src/sagemaker/serve/utils/optimize_utils.py

@@ -260,7 +260,7 @@ def _is_s3_uri(s3_uri: Optional[str]) -> bool:
 
 def _extract_optimization_config_and_env(
     quantization_config: Optional[Dict] = None, compilation_config: Optional[Dict] = None
-) -> Optional[Tuple[Optional[Dict], Optional[Dict]]]:
+) -> Optional[Tuple[Optional[Dict], Optional[Dict], Optional[Dict]]]:
     """Extracts optimization config and environment variables.
 
     Args:
@@ -272,25 +272,24 @@ def _extract_optimization_config_and_env(
             The optimization config and environment variables.
     """
     optimization_config = {}
-    quantization_override_env = {}
-    compilation_override_env = {}
+    quantization_override_env = (
+        quantization_config.get("OverrideEnvironment", {}) if quantization_config else None
+    )
+    compilation_override_env = (
+        compilation_config.get("OverrideEnvironment", {}) if compilation_config else None
+    )
 
-    if quantization_config:
+    if quantization_config is not None:
         optimization_config["ModelQuantizationConfig"] = quantization_config
-        quantization_override_env = quantization_config.get("OverrideEnvironment")
 
-    if compilation_config:
+    if compilation_config is not None:
         optimization_config["ModelCompilationConfig"] = compilation_config
-        compilation_override_env = compilation_config.get("OverrideEnvironment")
 
     # Return both dicts and environment variable if either is present
     if optimization_config:
-        return optimization_config, {
-            **(quantization_override_env or {}),
-            **(compilation_override_env or {}),
-        }
+        return optimization_config, quantization_override_env, compilation_override_env
 
-    return None, None
+    return None, None, None
 
 
 def _custom_speculative_decoding(

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

@@ -261,7 +261,7 @@ def test_is_s3_uri(s3_uri, expected):
 
 
 @pytest.mark.parametrize(
-    "quantization_config, compilation_config, expected_config, expected_env",
+    "quantization_config, compilation_config, expected_config, expected_quant_env, expected_compilation_env",
     [
         (
             None,
@@ -277,6 +277,7 @@ def test_is_s3_uri(s3_uri, expected):
                     }
                 },
             },
+            None,
             {
                 "OPTION_TENSOR_PARALLEL_DEGREE": "2",
             },
@@ -298,16 +299,22 @@ def test_is_s3_uri(s3_uri, expected):
             {
                 "OPTION_TENSOR_PARALLEL_DEGREE": "2",
             },
+            None,
         ),
-        (None, None, None, None),
+        (None, None, None, None, None),
     ],
 )
 def test_extract_optimization_config_and_env(
-    quantization_config, compilation_config, expected_config, expected_env
+    quantization_config,
+    compilation_config,
+    expected_config,
+    expected_quant_env,
+    expected_compilation_env,
 ):
     assert _extract_optimization_config_and_env(quantization_config, compilation_config) == (
         expected_config,
-        expected_env,
+        expected_quant_env,
+        expected_compilation_env,
     )