Fix incorrect assignment of ModelCompilationConfig and add UT.

Author: Joseph Zhang

src/sagemaker/serve/builder/jumpstart_builder.py

@@ -733,11 +733,17 @@ def _optimize_for_jumpstart(
         if (
             not optimization_config or not optimization_config.get("ModelCompilationConfig")
         ) and is_compilation:
-            override_env = override_env or pysdk_model_env_vars
+            # Ensure optimization_config exists
+            if not optimization_config:
+                optimization_config = {}
+
+            # Fallback to default if override_env is None or empty
+            if not override_env:
+                override_env = pysdk_model_env_vars
+
+            # Update optimization_config with ModelCompilationConfig
             optimization_config["ModelCompilationConfig"] = {
-                "ModelCompilationConfig": {
-                    "OverrideEnvironment": override_env,
-                }
+                "OverrideEnvironment": override_env,
             }
 
         if speculative_decoding_config:

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

@@ -1442,3 +1442,103 @@ def test_optimize_compile_for_jumpstart_with_neuron_env(
         self.assertEqual(optimized_model.env["OPTION_ROLLING_BATCH"], "auto")
         self.assertEqual(optimized_model.env["OPTION_MAX_ROLLING_BATCH_SIZE"], "4")
         self.assertEqual(optimized_model.env["OPTION_NEURON_OPTIMIZE_LEVEL"], "2")
+
+    @patch("sagemaker.serve.builder.jumpstart_builder._capture_telemetry", side_effect=None)
+    @patch.object(ModelBuilder, "_get_serve_setting", autospec=True)
+    @patch(
+        "sagemaker.serve.builder.jumpstart_builder.JumpStart._is_gated_model",
+        return_value=True,
+    )
+    @patch("sagemaker.serve.builder.jumpstart_builder.JumpStartModel")
+    @patch(
+        "sagemaker.serve.builder.jumpstart_builder.JumpStart._is_jumpstart_model_id",
+        return_value=True,
+    )
+    @patch("sagemaker.serve.builder.jumpstart_builder.JumpStart._create_pre_trained_js_model")
+    @patch(
+        "sagemaker.serve.builder.jumpstart_builder.prepare_tgi_js_resources",
+        return_value=({"model_type": "t5", "n_head": 71}, True),
+    )
+    def test_optimize_compile_for_jumpstart_without_compilation_config(
+        self,
+        mock_prepare_for_tgi,
+        mock_pre_trained_model,
+        mock_is_jumpstart_model,
+        mock_js_model,
+        mock_is_gated_model,
+        mock_serve_settings,
+        mock_telemetry,
+    ):
+        mock_sagemaker_session = Mock()
+        mock_metadata_config = Mock()
+        mock_sagemaker_session.wait_for_optimization_job.side_effect = (
+            lambda *args: mock_optimization_job_response
+        )
+
+        mock_metadata_config.resolved_config = {
+            "supported_inference_instance_types": ["ml.inf2.48xlarge"],
+            "hosting_neuron_model_id": "huggingface-llmneuron-mistral-7b",
+        }
+
+        mock_js_model.return_value = MagicMock()
+        mock_js_model.return_value.env = {
+            "SAGEMAKER_PROGRAM": "inference.py",
+            "ENDPOINT_SERVER_TIMEOUT": "3600",
+            "MODEL_CACHE_ROOT": "/opt/ml/model",
+            "SAGEMAKER_ENV": "1",
+            "HF_MODEL_ID": "/opt/ml/model",
+            "SAGEMAKER_MODEL_SERVER_WORKERS": "1",
+        }
+
+        mock_pre_trained_model.return_value = MagicMock()
+        mock_pre_trained_model.return_value.env = dict()
+        mock_pre_trained_model.return_value.config_name = "config_name"
+        mock_pre_trained_model.return_value.model_data = mock_model_data
+        mock_pre_trained_model.return_value.image_uri = mock_tgi_image_uri
+        mock_pre_trained_model.return_value.list_deployment_configs.return_value = (
+            DEPLOYMENT_CONFIGS
+        )
+        mock_pre_trained_model.return_value.deployment_config = DEPLOYMENT_CONFIGS[0]
+        mock_pre_trained_model.return_value._metadata_configs = {
+            "config_name": mock_metadata_config
+        }
+
+        sample_input = {
+            "inputs": "The diamondback terrapin or simply terrapin is a species "
+            "of turtle native to the brackish coastal tidal marshes of the",
+            "parameters": {"max_new_tokens": 1024},
+        }
+        sample_output = [
+            {
+                "generated_text": "The diamondback terrapin or simply terrapin is a "
+                "species of turtle native to the brackish coastal "
+                "tidal marshes of the east coast."
+            }
+        ]
+
+        model_builder = ModelBuilder(
+            model="meta-textgeneration-llama-3-70b",
+            schema_builder=SchemaBuilder(sample_input, sample_output),
+            sagemaker_session=mock_sagemaker_session,
+        )
+
+        optimized_model = model_builder.optimize(
+            accept_eula=True,
+            instance_type="ml.inf2.24xlarge",
+            output_path="s3://bucket/code/",
+        )
+
+        self.assertEqual(
+            optimized_model.image_uri,
+            mock_optimization_job_response["OptimizationOutput"]["RecommendedInferenceImage"],
+        )
+        self.assertEqual(
+            optimized_model.model_data["S3DataSource"]["S3Uri"],
+            mock_optimization_job_response["OutputConfig"]["S3OutputLocation"],
+        )
+        self.assertEqual(optimized_model.env["SAGEMAKER_PROGRAM"], "inference.py")
+        self.assertEqual(optimized_model.env["ENDPOINT_SERVER_TIMEOUT"], "3600")
+        self.assertEqual(optimized_model.env["MODEL_CACHE_ROOT"], "/opt/ml/model")
+        self.assertEqual(optimized_model.env["SAGEMAKER_ENV"], "1")
+        self.assertEqual(optimized_model.env["HF_MODEL_ID"], "/opt/ml/model")
+        self.assertEqual(optimized_model.env["SAGEMAKER_MODEL_SERVER_WORKERS"], "1")