chore: require config name and instance type in set_deployment_config (aws#4625)

* require config_name and instance_type in set config

* docstring

* add supported instance types check

* add more tests

* format

* fix tests

Author: Captainia

src/sagemaker/jumpstart/factory/model.py

@@ -544,7 +544,11 @@ def _add_resources_to_kwargs(kwargs: JumpStartModelInitKwargs) -> JumpStartModel
 
 
 def _add_config_name_to_kwargs(kwargs: JumpStartModelInitKwargs) -> JumpStartModelInitKwargs:
-    """Sets default config name to the kwargs. Returns full kwargs."""
+    """Sets default config name to the kwargs. Returns full kwargs.
+
+    Raises:
+        ValueError: If the instance_type is not supported with the current config.
+    """
 
     specs = verify_model_region_and_return_specs(
         model_id=kwargs.model_id,
@@ -565,6 +569,22 @@ def _add_config_name_to_kwargs(kwargs: JumpStartModelInitKwargs) -> JumpStartMod
             kwargs.config_name or specs.inference_configs.get_top_config_from_ranking().config_name
         )
 
+        if not kwargs.config_name:
+            return kwargs
+
+        if kwargs.config_name not in set(specs.inference_configs.configs.keys()):
+            raise ValueError(
+                f"Config {kwargs.config_name} is not supported for model {kwargs.model_id}."
+            )
+
+        resolved_config = specs.inference_configs.configs[kwargs.config_name].resolved_config
+        supported_instance_types = resolved_config.get("supported_inference_instance_types", [])
+        if kwargs.instance_type not in supported_instance_types:
+            raise ValueError(
+                f"Instance type {kwargs.instance_type} "
+                f"is not supported for config {kwargs.config_name}."
+            )
+
     return kwargs
 
 

src/sagemaker/jumpstart/model.py

@@ -429,16 +429,22 @@ def retrieve_example_payload(self) -> JumpStartSerializablePayload:
             sagemaker_session=self.sagemaker_session,
         )
 
-    def set_deployment_config(self, config_name: Optional[str]) -> None:
+    def set_deployment_config(self, config_name: str, instance_type: str) -> None:
         """Sets the deployment config to apply to the model.
 
         Args:
-            config_name (Optional[str]):
-                The name of the deployment config. Set to None to unset
-                any existing config that is applied to the model.
+            config_name (str):
+                The name of the deployment config to apply to the model.
+                Call list_deployment_configs to see the list of config names.
+            instance_type (str):
+                The instance_type that the model will use after setting
+                the config.
         """
         self.__init__(
-            model_id=self.model_id, model_version=self.model_version, config_name=config_name
+            model_id=self.model_id,
+            model_version=self.model_version,
+            instance_type=instance_type,
+            config_name=config_name,
         )
 
     @property

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

@@ -1614,7 +1614,25 @@ def test_model_set_deployment_config(
         mock_get_model_specs.reset_mock()
         mock_model_deploy.reset_mock()
         mock_get_model_specs.side_effect = get_prototype_spec_with_configs
-        model.set_deployment_config("neuron-inference")
+        model.set_deployment_config("neuron-inference", "ml.inf2.2xlarge")
+
+        assert model.config_name == "neuron-inference"
+
+        model.deploy()
+
+        mock_model_deploy.assert_called_once_with(
+            initial_instance_count=1,
+            instance_type="ml.inf2.2xlarge",
+            tags=[
+                {"Key": JumpStartTag.MODEL_ID, "Value": "pytorch-eqa-bert-base-cased"},
+                {"Key": JumpStartTag.MODEL_VERSION, "Value": "1.0.0"},
+                {"Key": JumpStartTag.MODEL_CONFIG_NAME, "Value": "neuron-inference"},
+            ],
+            wait=True,
+            endpoint_logging=False,
+        )
+        mock_model_deploy.reset_mock()
+        model.set_deployment_config("neuron-inference", "ml.inf2.xlarge")
 
         assert model.config_name == "neuron-inference"
 
@@ -1640,15 +1658,15 @@ def test_model_set_deployment_config(
     @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_unset_deployment_config(
+    def test_model_set_deployment_config_incompatible_instance_type_or_name(
         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_model_specs.side_effect = get_prototype_model_spec
         mock_get_manifest.side_effect = (
             lambda region, model_type, *args, **kwargs: get_prototype_manifest(region, model_type)
         )
@@ -1658,40 +1676,40 @@ def test_model_unset_deployment_config(
 
         mock_session.return_value = sagemaker_session
 
-        model = JumpStartModel(model_id=model_id, config_name="neuron-inference")
+        model = JumpStartModel(model_id=model_id)
 
-        assert model.config_name == "neuron-inference"
+        assert model.config_name is None
 
         model.deploy()
 
         mock_model_deploy.assert_called_once_with(
             initial_instance_count=1,
-            instance_type="ml.inf2.xlarge",
+            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.MODEL_CONFIG_NAME, "Value": "neuron-inference"},
             ],
             wait=True,
             endpoint_logging=False,
         )
 
         mock_get_model_specs.reset_mock()
         mock_model_deploy.reset_mock()
-        mock_get_model_specs.side_effect = get_prototype_model_spec
-        model.set_deployment_config(None)
-
-        model.deploy()
+        mock_get_model_specs.side_effect = get_prototype_spec_with_configs
+        with pytest.raises(ValueError) as error:
+            model.set_deployment_config("neuron-inference", "ml.inf2.32xlarge")
+        assert (
+            "Instance type ml.inf2.32xlarge is not supported for config neuron-inference."
+            in str(error)
+        )
 
-        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"},
-            ],
-            wait=True,
-            endpoint_logging=False,
+        with pytest.raises(ValueError) as error:
+            model.set_deployment_config("neuron-inference-unknown-name", "ml.inf2.32xlarge")
+        assert (
+            "Cannot find Jumpstart config name neuron-inference-unknown-name. "
+            "List of config names that is supported by the model: "
+            "['neuron-inference', 'neuron-inference-budget', 'gpu-inference-budget', 'gpu-inference']"
+            in str(error)
         )
 
     @mock.patch("sagemaker.jumpstart.model.get_init_kwargs")
@@ -1813,6 +1831,7 @@ def test_model_retrieve_deployment_config(
 
         expected = get_base_deployment_configs()[0]
         config_name = expected.get("DeploymentConfigName")
+        instance_type = expected.get("InstanceType")
         mock_get_init_kwargs.side_effect = lambda *args, **kwargs: get_mock_init_kwargs(
             model_id, config_name
         )
@@ -1821,17 +1840,13 @@ def test_model_retrieve_deployment_config(
 
         model = JumpStartModel(model_id=model_id)
 
-        model.set_deployment_config(config_name)
+        model.set_deployment_config(config_name, instance_type)
 
         self.assertEqual(model.deployment_config, expected)
 
         mock_get_init_kwargs.reset_mock()
         mock_get_init_kwargs.side_effect = lambda *args, **kwargs: get_mock_init_kwargs(model_id)
 
-        # Unset
-        model.set_deployment_config(None)
-        self.assertIsNone(model.deployment_config)
-
     @mock.patch("sagemaker.jumpstart.model.get_init_kwargs")
     @mock.patch("sagemaker.jumpstart.utils.verify_model_region_and_return_specs")
     @mock.patch("sagemaker.jumpstart.model.get_instance_rate_per_hour")