feat: update alt config to work with model packages

src/sagemaker/jumpstart/artifacts/model_packages.py

@@ -96,7 +96,10 @@ def _retrieve_model_package_arn(
         if instance_specific_arn is not None:
             return instance_specific_arn
 
-        if model_specs.hosting_model_package_arns is None:
+        if (
+            model_specs.hosting_model_package_arns is None
+            or model_specs.hosting_model_package_arns == {}
+        ):
             return None
 
         regional_arn = model_specs.hosting_model_package_arns.get(region)

src/sagemaker/jumpstart/types.py

@@ -951,7 +951,9 @@ 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_arns: Optional[Dict] = json_obj.get("hosting_model_package_arns")
+        self.hosting_model_package_arns: Optional[Dict] = json_obj.get(
+            "hosting_model_package_arns", {}
+        )
         self.hosting_use_script_uri: bool = json_obj.get("hosting_use_script_uri", True)
 
         self.hosting_instance_type_variants: Optional[JumpStartInstanceTypeVariants] = (

tests/unit/sagemaker/jumpstart/constants.py

@@ -7692,6 +7692,14 @@
             },
             "component_names": ["gpu-inference"],
         },
+        "gpu-inference-model-package": {
+            "benchmark_metrics": {
+                "ml.p3.2xlarge": [
+                    {"name": "Latency", "value": "100", "unit": "Tokens/S", "concurrency": 1}
+                ]
+            },
+            "component_names": ["gpu-inference-model-package"],
+        },
     },
     "inference_config_components": {
         "neuron-base": {
@@ -7733,6 +7741,15 @@
                 },
             },
         },
+        "gpu-inference-model-package": {
+            "default_inference_instance_type": "ml.p2.xlarge",
+            "supported_inference_instance_types": ["ml.p2.xlarge", "ml.p3.2xlarge"],
+            "hosting_model_package_arns": {
+                "us-west-2": "arn:aws:sagemaker:us-west-2:594846645681:model-package/ll"
+                "ama2-7b-v3-740347e540da35b4ab9f6fc0ab3fed2c"
+            },
+            "inference_environment_variables": [],
+        },
         "gpu-inference-budget": {
             "supported_inference_instance_types": ["ml.p2.xlarge", "ml.p3.2xlarge"],
             "hosting_artifact_key": "artifacts/meta-textgeneration-llama-2-7b/gpu-inference-budget/model/",

tests/unit/sagemaker/jumpstart/model/test_model.py

@@ -1651,6 +1651,74 @@ def test_model_set_deployment_config(
             endpoint_logging=False,
         )
 
+    @mock.patch(
+        "sagemaker.jumpstart.model.get_jumpstart_configs", side_effect=lambda *args, **kwargs: {}
+    )
+    @mock.patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor._get_manifest")
+    @mock.patch("sagemaker.jumpstart.factory.model.Session")
+    @mock.patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
+    @mock.patch("sagemaker.jumpstart.model.Model.deploy")
+    @mock.patch("sagemaker.jumpstart.factory.model.JUMPSTART_DEFAULT_REGION_NAME", region)
+    def test_model_set_deployment_config_model_package(
+        self,
+        mock_model_deploy: mock.Mock,
+        mock_get_model_specs: mock.Mock,
+        mock_session: mock.Mock,
+        mock_get_manifest: mock.Mock,
+        mock_get_jumpstart_configs: mock.Mock,
+    ):
+        mock_get_model_specs.side_effect = get_prototype_spec_with_configs
+        mock_get_manifest.side_effect = (
+            lambda region, model_type, *args, **kwargs: get_prototype_manifest(region, model_type)
+        )
+        mock_model_deploy.return_value = default_predictor
+
+        model_id, _ = "pytorch-eqa-bert-base-cased", "*"
+
+        mock_session.return_value = sagemaker_session
+
+        model = JumpStartModel(model_id=model_id)
+
+        assert model.config_name == "neuron-inference"
+
+        model.deploy()
+
+        mock_model_deploy.assert_called_once_with(
+            initial_instance_count=1,
+            instance_type="ml.inf2.xlarge",
+            tags=[
+                {"Key": JumpStartTag.MODEL_ID, "Value": "pytorch-eqa-bert-base-cased"},
+                {"Key": JumpStartTag.MODEL_VERSION, "Value": "1.0.0"},
+                {"Key": JumpStartTag.INFERENCE_CONFIG_NAME, "Value": "neuron-inference"},
+            ],
+            wait=True,
+            endpoint_logging=False,
+        )
+
+        mock_model_deploy.reset_mock()
+
+        model.set_deployment_config(
+            config_name="gpu-inference-model-package", instance_type="ml.p2.xlarge"
+        )
+
+        assert (
+            model.model_package_arn
+            == "arn:aws:sagemaker:us-west-2:594846645681:model-package/llama2-7b-v3-740347e540da35b4ab9f6fc0ab3fed2c"
+        )
+        model.deploy()
+
+        mock_model_deploy.assert_called_once_with(
+            initial_instance_count=1,
+            instance_type="ml.p2.xlarge",
+            tags=[
+                {"Key": JumpStartTag.MODEL_ID, "Value": "pytorch-eqa-bert-base-cased"},
+                {"Key": JumpStartTag.MODEL_VERSION, "Value": "1.0.0"},
+                {"Key": JumpStartTag.INFERENCE_CONFIG_NAME, "Value": "gpu-inference-model-package"},
+            ],
+            wait=True,
+            endpoint_logging=False,
+        )
+
     @mock.patch(
         "sagemaker.jumpstart.model.get_jumpstart_configs", side_effect=lambda *args, **kwargs: {}
     )