chore: use hosting_model_package_arns

Author: evakravi

src/sagemaker/jumpstart/artifacts/model_packages.py

@@ -50,7 +50,7 @@ def _retrieve_model_package_arn(
             an exception if the version of the model is deprecated. (Default: False).
 
     Returns:
-        list: the model package arn to use for the model or None.
+        str: the model package arn to use for the model or None.
     """
 
     if region is None:
@@ -66,6 +66,12 @@ def _retrieve_model_package_arn(
     )
 
     if scope == JumpStartScriptScope.INFERENCE:
-        return model_specs.hosting_model_package_arn
+
+        if model_specs.hosting_model_package_arns is None:
+            return None
+
+        regional_arn = model_specs.hosting_model_package_arns.get(region)
+
+        return regional_arn
 
     raise NotImplementedError(f"Model Package ARN not supported for scope: '{scope}'")

src/sagemaker/jumpstart/types.py

@@ -352,7 +352,7 @@ class JumpStartModelSpecs(JumpStartDataHolderType):
         "training_enable_network_isolation",
         "resource_name_base",
         "hosting_eula_key",
-        "hosting_model_package_arn",
+        "hosting_model_package_arns",
     ]
 
     def __init__(self, spec: Dict[str, Any]):
@@ -423,7 +423,7 @@ def from_json(self, json_obj: Dict[str, Any]) -> None:
 
         self.hosting_eula_key: Optional[str] = json_obj.get("hosting_eula_key")
 
-        self.hosting_model_package_arn: Optional[str] = json_obj.get("hosting_model_package_arn")
+        self.hosting_model_package_arns: Optional[Dict] = json_obj.get("hosting_model_package_arns")
 
         if self.training_supported:
             self.training_ecr_specs: JumpStartECRSpecs = JumpStartECRSpecs(

tests/unit/sagemaker/jumpstart/constants.py

@@ -2299,7 +2299,7 @@
     "training_script_key": "source-directory-tarballs/pytorch/transfer_learning/ic/v1.0.0/sourcedir.tar.gz",
     "training_prepacked_script_key": None,
     "hosting_prepacked_artifact_key": None,
-    "hosting_model_package_arn": None,
+    "hosting_model_package_arns": None,
     "hosting_eula_key": False,
     "hyperparameters": [
         {