fix: config logic support

Author: tanvikab

src/sagemaker/jumpstart/factory/estimator.py

@@ -837,6 +837,7 @@ def _add_config_name_to_kwargs(
 
     kwargs.config_name = kwargs.config_name or get_top_ranked_config_name(
         scope=JumpStartScriptScope.TRAINING,
+        instance_type=kwargs.instance_type,
         **get_model_info_default_kwargs(kwargs, include_config_name=False),
     )
 

src/sagemaker/jumpstart/factory/model.py

@@ -559,6 +559,7 @@ def _add_config_name_to_init_kwargs(kwargs: JumpStartModelInitKwargs) -> JumpSta
     kwargs.config_name = kwargs.config_name or get_top_ranked_config_name(
         **get_model_info_default_kwargs(kwargs, include_config_name=False),
         scope=JumpStartScriptScope.INFERENCE,
+        instance_type=kwargs.instance_type,
     )
 
     if kwargs.config_name is None:
@@ -618,6 +619,7 @@ def _add_config_name_to_deploy_kwargs(
         default_config_name = kwargs.config_name or get_top_ranked_config_name(
             **get_model_info_default_kwargs(kwargs, include_config_name=False),
             scope=JumpStartScriptScope.INFERENCE,
+            instance_type=kwargs.instance_type,
         )
 
     kwargs.config_name = kwargs.config_name or default_config_name
@@ -927,6 +929,12 @@ def get_init_kwargs(
 
     model_init_kwargs = _add_vulnerable_and_deprecated_status_to_kwargs(kwargs=model_init_kwargs)
     model_init_kwargs = _add_model_version_to_kwargs(kwargs=model_init_kwargs)
+
+    # Add instance type before config selection so config compatibility can be checked
+    model_init_kwargs = _add_instance_type_to_kwargs(
+        kwargs=model_init_kwargs, disable_instance_type_logging=disable_instance_type_logging
+    )
+
     model_init_kwargs = _add_config_name_to_init_kwargs(kwargs=model_init_kwargs)
 
     model_init_kwargs = _add_sagemaker_session_with_custom_user_agent_to_kwargs(
@@ -936,10 +944,6 @@ def get_init_kwargs(
 
     model_init_kwargs = _add_model_name_to_kwargs(kwargs=model_init_kwargs)
 
-    model_init_kwargs = _add_instance_type_to_kwargs(
-        kwargs=model_init_kwargs, disable_instance_type_logging=disable_instance_type_logging
-    )
-
     model_init_kwargs = _add_image_uri_to_kwargs(kwargs=model_init_kwargs)
 
     if hub_arn:

src/sagemaker/jumpstart/types.py

@@ -1723,10 +1723,10 @@ def get_top_config_from_ranking(
             ranked_config_names = rankings.rankings
         for config_name in ranked_config_names:
             resolved_config = self.configs[config_name].resolved_config
-            if instance_type and instance_type not in getattr(
-                resolved_config, instance_type_attribute
-            ):
-                continue
+            if instance_type:
+                supported_instance_types = getattr(resolved_config, instance_type_attribute, [])
+                if supported_instance_types and instance_type not in supported_instance_types:
+                    continue
             return self.configs[config_name]
 
         return None

src/sagemaker/jumpstart/utils.py

@@ -1233,9 +1233,14 @@ def get_top_ranked_config_name(
     tolerate_vulnerable_model: bool = False,
     hub_arn: Optional[str] = None,
     ranking_name: enums.JumpStartConfigRankingName = enums.JumpStartConfigRankingName.DEFAULT,
+    instance_type: Optional[str] = None,
 ) -> Optional[str]:
     """Returns the top ranked config name for the given model ID and region.
 
+    Args:
+        instance_type (Optional[str]): The instance type to filter configs by compatibility.
+            If provided, only configs that support this instance type will be considered.
+
     Raises:
         ValueError: If the script scope is not supported by JumpStart.
     """
@@ -1254,15 +1259,15 @@ def get_top_ranked_config_name(
     if scope == enums.JumpStartScriptScope.INFERENCE:
         return (
             model_specs.inference_configs.get_top_config_from_ranking(
-                ranking_name=ranking_name
+                ranking_name=ranking_name, instance_type=instance_type
             ).config_name
             if model_specs.inference_configs
             else None
         )
     if scope == enums.JumpStartScriptScope.TRAINING:
         return (
             model_specs.training_configs.get_top_config_from_ranking(
-                ranking_name=ranking_name
+                ranking_name=ranking_name, instance_type=instance_type
             ).config_name
             if model_specs.training_configs
             else None