fix: Remove hardcoded lmi config name in ModelBuilder optimize

Replace hardcoded config_name='lmi' in ModelBuilder.optimize() with
the model's dynamically resolved default config name. The JumpStart
metadata for models like llama-3-1-8b-instruct no longer includes an
'lmi' config, causing ValueError during optimization.

Also fix test_jumpstart_session_with_config_name to use the model's
resolved config_name instead of hardcoding 'tgi', and update unit
test assertions in test_js_builder.py accordingly.

Author: nargokul

src/sagemaker/serve/builder/model_builder.py

@@ -1728,9 +1728,11 @@ def _model_builder_optimize_wrapper(
         if self._is_jumpstart_model_id():
             self.build(mode=self.mode, sagemaker_session=self.sagemaker_session)
             if self.pysdk_model:
-                self.pysdk_model.set_deployment_config(
-                    instance_type=instance_type, config_name="lmi"
-                )
+                config_name = self.pysdk_model.config_name
+                if config_name:
+                    self.pysdk_model.set_deployment_config(
+                        instance_type=instance_type, config_name=config_name
+                    )
             input_args = self._optimize_for_jumpstart(
                 output_path=output_path,
                 instance_type=instance_type,

tests/integ/sagemaker/jumpstart/model/test_jumpstart_model.py

@@ -414,7 +414,7 @@ def test_jumpstart_session_with_config_name():
             pass
 
     assert (
-        "md/js_model_id#meta-textgeneration-llama-2-7b md/js_model_ver#* md/js_config#tgi"
+        f"md/js_model_id#meta-textgeneration-llama-2-7b md/js_model_ver#* md/js_config#{model.config_name}"
         in mock_make_request.call_args[0][1]["headers"]["User-Agent"]
     )
 

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

@@ -1696,7 +1696,7 @@ def test_optimize_on_js_model_should_ignore_pre_optimized_configurations(
 
         assert mock_lmi_js_model.set_deployment_config.call_args_list[0].kwargs == {
             "instance_type": "ml.g5.24xlarge",
-            "config_name": "lmi",
+            "config_name": mock_lmi_js_model.config_name,
         }
         assert optimized_model.env == {
             "SAGEMAKER_PROGRAM": "inference.py",
@@ -1784,7 +1784,7 @@ def test_optimize_on_js_model_should_ignore_pre_optimized_configurations_no_over
 
         assert mock_lmi_js_model.set_deployment_config.call_args_list[0].kwargs == {
             "instance_type": "ml.g5.24xlarge",
-            "config_name": "lmi",
+            "config_name": mock_lmi_js_model.config_name,
         }
         assert optimized_model.env == {
             "SAGEMAKER_PROGRAM": "inference.py",