add UTs for JumpStart deployment

Author: gwang111

src/sagemaker/jumpstart/model.py

@@ -817,12 +817,14 @@ def deploy(
                 f"{EndpointType.INFERENCE_COMPONENT_BASED} is not supported for Proprietary models."
             )
 
+        print(self.additional_model_data_sources)
         self.additional_model_data_sources = _add_model_access_configs_to_model_data_sources(
             self.additional_model_data_sources,
             deploy_kwargs.model_access_configs,
             deploy_kwargs.model_id,
             deploy_kwargs.region,
         )
+        print(self.additional_model_data_sources)
 
         try:
             predictor = super(JumpStartModel, self).deploy(**deploy_kwargs.to_kwargs_dict())

src/sagemaker/jumpstart/utils.py

@@ -1544,15 +1544,21 @@ def _add_model_access_configs_to_model_data_sources(
     region: str,
 ):
     """Sets AcceptEula to True for gated speculative decoding models"""
+    print(model_data_sources)
 
     if not model_data_sources:
         return model_data_sources
 
     acked_model_data_sources = []
     for model_data_source in model_data_sources:
         hosting_eula_key = model_data_source.get("HostingEulaKey")
+        mutable_model_data_source = model_data_source.copy()
         if hosting_eula_key:
-            if not model_access_configs or not model_access_configs.get(model_id):
+            if (
+                not model_access_configs
+                or not model_access_configs.get(model_id)
+                or not model_access_configs.get(model_id).accept_eula
+            ):
                 eula_message_template = (
                     "{model_source}{base_eula_message}{model_access_configs_message}"
                 )
@@ -1572,14 +1578,14 @@ def _add_model_access_configs_to_model_data_sources(
                         ),
                     )
                 )
-            acked_model_data_source = model_data_source.copy()
-            acked_model_data_source.pop("HostingEulaKey")
-            acked_model_data_source["S3DataSource"]["ModelAccessConfig"] = (
+            mutable_model_data_source.pop("HostingEulaKey")  # pop when model access config is applied
+            mutable_model_data_source["S3DataSource"]["ModelAccessConfig"] = (
                 camel_case_to_pascal_case(model_access_configs.get(model_id).model_dump())
             )
-            acked_model_data_sources.append(acked_model_data_source)
+            acked_model_data_sources.append(mutable_model_data_source)
         else:
-            acked_model_data_sources.append(model_data_source)
+            mutable_model_data_source.pop("HostingEulaKey")  # pop when model access config is not applicable
+            acked_model_data_sources.append(mutable_model_data_source)
     return acked_model_data_sources
 
 

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

@@ -19,6 +19,7 @@
 import pandas as pd
 from mock import MagicMock, Mock
 import pytest
+from sagemaker_core.shapes import ModelAccessConfig
 from sagemaker.async_inference.async_inference_config import AsyncInferenceConfig
 from sagemaker.jumpstart.artifacts.environment_variables import (
     _retrieve_default_environment_variables,
@@ -54,6 +55,7 @@
     get_base_deployment_configs,
     get_base_spec_with_prototype_configs_with_missing_benchmarks,
     append_instance_stat_metrics,
+    append_gated_draft_model_specs_to_jumpstart_model_spec,
 )
 import boto3
 
@@ -772,6 +774,7 @@ def test_jumpstart_model_kwargs_match_parent_class(self):
 
         init_args_to_skip: Set[str] = set(["model_reference_arn"])
         deploy_args_to_skip: Set[str] = set(["kwargs", "model_reference_arn"])
+        deploy_args_removed_at_deploy_time: Set[str] = set(["model_access_configs"])
 
         parent_class_init = Model.__init__
         parent_class_init_args = set(signature(parent_class_init).parameters.keys())
@@ -798,8 +801,9 @@ def test_jumpstart_model_kwargs_match_parent_class(self):
         js_class_deploy = JumpStartModel.deploy
         js_class_deploy_args = set(signature(js_class_deploy).parameters.keys())
 
-        assert js_class_deploy_args - parent_class_deploy_args == set()
-        assert parent_class_deploy_args - js_class_deploy_args == deploy_args_to_skip
+        assert js_class_deploy_args - parent_class_deploy_args - deploy_args_removed_at_deploy_time == set()
+        assert (parent_class_deploy_args - js_class_deploy_args - deploy_args_removed_at_deploy_time ==
+                deploy_args_to_skip)
 
     @mock.patch(
         "sagemaker.jumpstart.model.get_jumpstart_configs", side_effect=lambda *args, **kwargs: {}
@@ -1762,6 +1766,89 @@ 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_and_deploy_for_gated_draft_model(
+            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,
+    ):
+        # WHERE
+        mock_get_model_specs.side_effect = append_gated_draft_model_specs_to_jumpstart_model_spec
+        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 is None
+
+        # WHEN
+        model.deploy(model_access_configs={"pytorch-eqa-bert-base-cased":ModelAccessConfig(accept_eula=True)})
+
+        # THEN
+        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,
+        )
+
+    @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_and_deploy_for_gated_draft_model_no_model_access_configs(
+            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,
+    ):
+        # WHERE
+        mock_get_model_specs.side_effect = append_gated_draft_model_specs_to_jumpstart_model_spec
+        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 is None
+
+        # WHEN / THEN
+        with self.assertRaises(ValueError):
+            model.deploy()
+
     @mock.patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor._get_manifest")
     @mock.patch(
         "sagemaker.jumpstart.factory.model.get_default_jumpstart_session_with_user_agent_suffix"
@@ -1810,6 +1897,7 @@ def test_model_deployment_config_additional_model_data_source(
                         "S3Uri": "s3://sagemaker-sd-models-prod-us-west-2/key/to/draft/model/artifact/",
                         "ModelAccessConfig": {"AcceptEula": False},
                     },
+                    "HostingEulaKey": None,
                 }
             ],
         )