fix: unit test for matching args

Author: evakravi

src/sagemaker/jumpstart/estimator.py

@@ -663,7 +663,7 @@ def deploy(
         role: Optional[str] = None,
         predictor_cls: Optional[callable] = None,
         env: Optional[Dict[str, Union[str, PipelineVariable]]] = None,
-        name: Optional[str] = None,
+        model_name: Optional[str] = None,
         vpc_config: Optional[Dict[str, List[Union[str, PipelineVariable]]]] = None,
         sagemaker_session: Optional[session.Session] = None,
         enable_network_isolation: Union[bool, PipelineVariable] = None,
@@ -675,6 +675,7 @@ def deploy(
         container_log_level: Optional[Union[int, PipelineVariable]] = None,
         dependencies: Optional[List[str]] = None,
         git_config: Optional[Dict[str, str]] = None,
+        use_compiled_model: bool = False,
     ) -> PredictorBase:
         """Creates endpoint from training job.
 
@@ -766,7 +767,7 @@ def deploy(
                 function on the created endpoint name. (Default: None).
             env (Optional[dict[str, str] or dict[str, PipelineVariable]]): Environment variables
                 to run with ``image_uri`` when hosted in SageMaker. (Default: None).
-            name (Optional[str]): The model name. If None, a default model name will be
+            model_name (Optional[str]): The model name. If None, a default model name will be
                 selected on each ``deploy``. (Default: None).
             vpc_config (Optional[Union[dict[str, list[str]],dict[str, list[PipelineVariable]]]]):
                 The VpcConfig set on the model (Default: None)
@@ -909,6 +910,8 @@ def deploy(
                 the SageMaker Python SDK attempts to use either the CodeCommit
                 credential helper or local credential storage for authentication.
                 (Default: None).
+            use_compiled_model (bool): Flag to select whether to use compiled
+                (optimized) model. (Default: False).
         """
 
         self.orig_predictor_cls = predictor_cls
@@ -948,7 +951,7 @@ def deploy(
             role=role,
             predictor_cls=predictor_cls,
             env=env,
-            name=name,
+            model_name=model_name,
             vpc_config=vpc_config,
             sagemaker_session=sagemaker_session,
             enable_network_isolation=enable_network_isolation,
@@ -960,6 +963,7 @@ def deploy(
             container_log_level=container_log_level,
             dependencies=dependencies,
             git_config=git_config,
+            use_compiled_model=use_compiled_model,
         )
 
         predictor = super(JumpStartEstimator, self).deploy(

src/sagemaker/jumpstart/factory/estimator.py

@@ -247,7 +247,6 @@ def get_deploy_kwargs(
     role: Optional[str] = None,
     predictor_cls: Optional[callable] = None,
     env: Optional[Dict[str, Union[str, PipelineVariable]]] = None,
-    name: Optional[str] = None,
     vpc_config: Optional[Dict[str, List[Union[str, PipelineVariable]]]] = None,
     sagemaker_session: Optional[Session] = None,
     enable_network_isolation: Union[bool, PipelineVariable] = None,
@@ -261,6 +260,8 @@ def get_deploy_kwargs(
     git_config: Optional[Dict[str, str]] = None,
     tolerate_deprecated_model: Optional[bool] = None,
     tolerate_vulnerable_model: Optional[bool] = None,
+    use_compiled_model: Optional[bool] = None,
+    model_name: Optional[str] = None,
 ) -> JumpStartEstimatorDeployKwargs:
     """Returns kwargs required to call `deploy` on `sagemaker.estimator.Estimator` object."""
 
@@ -301,7 +302,7 @@ def get_deploy_kwargs(
         env=env,
         predictor_cls=predictor_cls,
         role=role,
-        name=name,
+        name=model_name,
         vpc_config=vpc_config,
         sagemaker_session=sagemaker_session,
         enable_network_isolation=enable_network_isolation,
@@ -344,7 +345,7 @@ def get_deploy_kwargs(
         inference_recommendation_id=model_deploy_kwargs.inference_recommendation_id,
         explainer_config=model_deploy_kwargs.explainer_config,
         role=model_init_kwargs.role,
-        name=model_init_kwargs.name,
+        model_name=model_init_kwargs.name,
         vpc_config=model_init_kwargs.vpc_config,
         sagemaker_session=model_init_kwargs.sagemaker_session,
         enable_network_isolation=model_init_kwargs.enable_network_isolation,
@@ -356,6 +357,7 @@ def get_deploy_kwargs(
         git_config=model_init_kwargs.git_config,
         tolerate_vulnerable_model=model_init_kwargs.tolerate_vulnerable_model,
         tolerate_deprecated_model=model_init_kwargs.tolerate_deprecated_model,
+        use_compiled_model=use_compiled_model,
     )
 
     return estimator_deploy_kwargs

src/sagemaker/jumpstart/types.py

@@ -974,7 +974,6 @@ class JumpStartEstimatorDeployKwargs(JumpStartKwargs):
         "inference_recommendation_id",
         "explainer_config",
         "role",
-        "name",
         "vpc_config",
         "sagemaker_session",
         "enable_network_isolation",
@@ -986,6 +985,8 @@ class JumpStartEstimatorDeployKwargs(JumpStartKwargs):
         "git_config",
         "tolerate_deprecated_model",
         "tolerate_vulnerable_model",
+        "model_name",
+        "use_compiled_model",
     ]
 
     SERIALIZATION_EXCLUSION_SET = {
@@ -1023,7 +1024,7 @@ def __init__(
         role: Optional[str] = None,
         predictor_cls: Optional[callable] = None,
         env: Optional[Dict[str, Union[str, Any]]] = None,
-        name: Optional[str] = None,
+        model_name: Optional[str] = None,
         vpc_config: Optional[Dict[str, List[Union[str, Any]]]] = None,
         sagemaker_session: Optional[Any] = None,
         enable_network_isolation: Union[bool, Any] = None,
@@ -1037,6 +1038,7 @@ def __init__(
         git_config: Optional[Dict[str, str]] = None,
         tolerate_deprecated_model: Optional[bool] = None,
         tolerate_vulnerable_model: Optional[bool] = None,
+        use_compiled_model: bool = False,
     ) -> None:
         """Instantiates JumpStartEstimatorInitKwargs object."""
 
@@ -1066,7 +1068,7 @@ def __init__(
         self.inference_recommendation_id = inference_recommendation_id
         self.explainer_config = explainer_config
         self.role = role
-        self.name = name
+        self.model_name = model_name
         self.vpc_config = vpc_config
         self.sagemaker_session = sagemaker_session
         self.enable_network_isolation = enable_network_isolation
@@ -1078,3 +1080,4 @@ def __init__(
         self.git_config = git_config
         self.tolerate_deprecated_model = tolerate_deprecated_model
         self.tolerate_vulnerable_model = tolerate_vulnerable_model
+        self.use_compiled_model = use_compiled_model

tests/unit/sagemaker/jumpstart/estimator/test_estimator.py

@@ -12,7 +12,7 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 import time
-from typing import Optional
+from typing import Optional, Set
 from unittest import mock
 import unittest
 from inspect import signature
@@ -142,8 +142,9 @@ def test_non_prepacked(
             predictor_cls=Predictor,
             role=execution_role,
             wait=True,
+            use_compiled_model=False,
             enable_network_isolation=False,
-            name="blahblahblah-9876",
+            model_name="blahblahblah-9876",
             endpoint_name="blahblahblah-9876",
         )
 
@@ -234,6 +235,7 @@ def test_prepacked(
             predictor_cls=Predictor,
             role=execution_role,
             wait=True,
+            use_compiled_model=False,
             enable_network_isolation=False,
         )
 
@@ -503,7 +505,8 @@ def evaluate_estimator_workflow_with_kwargs(
                 "predictor_cls": Predictor,
                 "role": init_kwargs["role"],
                 "enable_network_isolation": False,
-                "name": "blahblahblah-1234",
+                "use_compiled_model": False,
+                "model_name": "blahblahblah-1234",
                 "endpoint_name": "blahblahblah-1234",
             },
             deploy_kwargs,
@@ -512,6 +515,15 @@ def evaluate_estimator_workflow_with_kwargs(
         mock_estimator_deploy.assert_called_once_with(**expected_deploy_kwargs)
 
     def test_jumpstart_estimator_kwargs_match_parent_class(self):
+
+        """If you add arguments to <Estimator constructor>, this test will fail.
+        Please add the new argument to the skip set below,
+        and cut a ticket sev-3 to JumpStart team: AWS > SageMaker > JumpStart"""
+
+        init_args_to_skip: Set[str] = set(["kwargs"])
+        fit_args_to_skip: Set[str] = set()
+        deploy_args_to_skip: Set[str] = set(["kwargs"])
+
         parent_class_init = Estimator.__init__
         parent_class_init_args = set(signature(parent_class_init).parameters.keys())
 
@@ -525,6 +537,7 @@ def test_jumpstart_estimator_kwargs_match_parent_class(self):
             "tolerate_vulnerable_model",
             "tolerate_deprecated_model",
         }
+        assert parent_class_init_args - js_class_init_args == init_args_to_skip
 
         parent_class_fit = Estimator.fit
         parent_class_fit_args = set(signature(parent_class_fit).parameters.keys())
@@ -533,6 +546,7 @@ def test_jumpstart_estimator_kwargs_match_parent_class(self):
         js_class_fit_args = set(signature(js_class_fit).parameters.keys())
 
         assert js_class_fit_args - parent_class_fit_args == set()
+        assert parent_class_fit_args - js_class_fit_args == fit_args_to_skip
 
         model_class_init = Model.__init__
         model_class_init_args = set(signature(model_class_init).parameters.keys())
@@ -546,7 +560,9 @@ def test_jumpstart_estimator_kwargs_match_parent_class(self):
         assert js_class_deploy_args - parent_class_deploy_args == model_class_init_args - {
             "model_data",
             "self",
+            "name",
         }
+        assert parent_class_deploy_args - js_class_deploy_args == deploy_args_to_skip
 
     @mock.patch("sagemaker.jumpstart.estimator.get_init_kwargs")
     @mock.patch("sagemaker.jumpstart.estimator.Estimator.__init__")
@@ -842,9 +858,10 @@ def test_training_passes_role_to_deploy(
             },
             predictor_cls=Predictor,
             wait=True,
+            use_compiled_model=False,
             role=mock_role,
             enable_network_isolation=False,
-            name="blahblahblah-3456",
+            model_name="blahblahblah-3456",
             endpoint_name="blahblahblah-3456",
         )
 
@@ -916,9 +933,10 @@ def test_training_passes_session_to_deploy(
             },
             predictor_cls=Predictor,
             wait=True,
+            use_compiled_model=False,
             role=mock_role,
             enable_network_isolation=False,
-            name="blahblahblah-3456",
+            model_name="blahblahblah-3456",
             endpoint_name="blahblahblah-3456",
         )
 

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

@@ -12,7 +12,7 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 from inspect import signature
-from typing import Optional
+from typing import Optional, Set
 from unittest import mock
 import unittest
 import pytest
@@ -322,6 +322,14 @@ def evaluate_model_workflow_with_kwargs(
         mock_model_deploy.assert_called_once_with(**expected_deploy_kwargs)
 
     def test_jumpstart_model_kwargs_match_parent_class(self):
+
+        """If you add arguments to <Model constructor>, this test will fail.
+        Please add the new argument to the skip set below,
+        and cut a ticket sev-3 to JumpStart team: AWS > SageMaker > JumpStart"""
+
+        init_args_to_skip: Set[str] = set()
+        deploy_args_to_skip: Set[str] = set(["kwargs"])
+
         parent_class_init = Model.__init__
         parent_class_init_args = set(signature(parent_class_init).parameters.keys())
 
@@ -336,6 +344,7 @@ def test_jumpstart_model_kwargs_match_parent_class(self):
             "tolerate_deprecated_model",
             "instance_type",
         }
+        assert parent_class_init_args - js_class_init_args == init_args_to_skip
 
         parent_class_deploy = Model.deploy
         parent_class_deploy_args = set(signature(parent_class_deploy).parameters.keys())
@@ -344,6 +353,7 @@ def test_jumpstart_model_kwargs_match_parent_class(self):
         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
 
     @mock.patch("sagemaker.jumpstart.model.get_init_kwargs")
     @mock.patch("sagemaker.jumpstart.model.Model.__init__")