fix: ModelReference deployment for Alt Configs models

Author: Malav Shastri

src/sagemaker/jumpstart/factory/estimator.py

@@ -912,6 +912,7 @@ def _add_config_name_to_kwargs(
         tolerate_deprecated_model=kwargs.tolerate_deprecated_model,
         sagemaker_session=kwargs.sagemaker_session,
         config_name=kwargs.config_name,
+        hub_arn=kwargs.hub_arn,
     )
 
     if specs.training_configs and specs.training_configs.get_top_config_from_ranking():

src/sagemaker/jumpstart/factory/model.py

@@ -672,9 +672,14 @@ def _add_config_name_to_init_kwargs(kwargs: JumpStartModelInitKwargs) -> JumpSta
         sagemaker_session=kwargs.sagemaker_session,
         model_type=kwargs.model_type,
         config_name=kwargs.config_name,
+        hub_arn=kwargs.hub_arn,
     )
     if specs.inference_configs:
-        default_config_name = specs.inference_configs.get_top_config_from_ranking().config_name
+        default_config_name = (
+            specs.inference_configs.get_top_config_from_ranking().config_name
+            if specs.inference_configs.get_top_config_from_ranking()
+            else None
+        )
         kwargs.config_name = kwargs.config_name or default_config_name
 
         if not kwargs.config_name:
@@ -707,6 +712,7 @@ def _add_additional_model_data_sources_to_kwargs(
         sagemaker_session=kwargs.sagemaker_session,
         model_type=kwargs.model_type,
         config_name=kwargs.config_name,
+        hub_arn=kwargs.hub_arn,
     )
     # Append speculative decoding data source from metadata
     speculative_decoding_data_sources = specs.get_speculative_decoding_s3_data_sources()
@@ -750,6 +756,7 @@ def _add_config_name_to_deploy_kwargs(
         sagemaker_session=kwargs.sagemaker_session,
         model_type=kwargs.model_type,
         config_name=kwargs.config_name,
+        hub_arn=kwargs.hub_arn,
     )
 
     if training_config_name:
@@ -759,7 +766,11 @@ def _add_config_name_to_deploy_kwargs(
         return kwargs
 
     if specs.inference_configs:
-        default_config_name = specs.inference_configs.get_top_config_from_ranking().config_name
+        default_config_name = (
+            specs.inference_configs.get_top_config_from_ranking().config_name
+            if specs.inference_configs.get_top_config_from_ranking()
+            else None
+        )
         kwargs.config_name = kwargs.config_name or default_config_name
 
     return kwargs

src/sagemaker/jumpstart/model.py

@@ -393,6 +393,7 @@ def _validate_model_id_and_type():
             model_version=self.model_version,
             sagemaker_session=self.sagemaker_session,
             model_type=self.model_type,
+            hub_arn=self.hub_arn,
         )
 
     def log_subscription_warning(self) -> None:

src/sagemaker/jumpstart/types.py

@@ -1732,32 +1732,15 @@ def from_json(self, json_obj: Dict[str, Any]) -> None:
         )
         self.inference_config_rankings: Optional[Dict[str, JumpStartConfigRanking]] = (
             {
-                alias: JumpStartConfigRanking(ranking)
+                alias: JumpStartConfigRanking(ranking, is_hub_content=self._is_hub_content)
                 for alias, ranking in json_obj["inference_config_rankings"].items()
             }
             if json_obj.get("inference_config_rankings")
             else None
         )
-        inference_configs_dict: Optional[Dict[str, JumpStartMetadataConfig]] = (
-            {
-                alias: JumpStartMetadataConfig(
-                    alias,
-                    config,
-                    json_obj,
-                    (
-                        {
-                            component_name: self.inference_config_components.get(component_name)
-                            for component_name in config.get("component_names")
-                        }
-                        if config and config.get("component_names")
-                        else None
-                    ),
-                )
-                for alias, config in json_obj["inference_configs"].items()
-            }
-            if json_obj.get("inference_configs")
-            else None
-        )
+        inference_configs_dict: Optional[Dict[str, JumpStartMetadataConfig]] = json_obj[
+            "inference_configs"
+        ]["Configs"]
         self.inference_configs: Optional[JumpStartMetadataConfigs] = (
             JumpStartMetadataConfigs(
                 inference_configs_dict,

src/sagemaker/jumpstart/utils.py

@@ -1074,6 +1074,7 @@ def get_jumpstart_configs(
     sagemaker_session: Optional[Session] = constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
     scope: enums.JumpStartScriptScope = enums.JumpStartScriptScope.INFERENCE,
     model_type: enums.JumpStartModelType = enums.JumpStartModelType.OPEN_WEIGHTS,
+    hub_arn: Optional[str] = None,
 ) -> Dict[str, JumpStartMetadataConfig]:
     """Returns metadata configs for the given model ID and region.
 
@@ -1087,6 +1088,7 @@ def get_jumpstart_configs(
         sagemaker_session=sagemaker_session,
         scope=scope,
         model_type=model_type,
+        hub_arn=hub_arn,
     )
 
     if scope == enums.JumpStartScriptScope.INFERENCE: