feature: inference recommendation id deployment integ test (aws#806)

* Add integ test and refactor rec id deploy

* Revert ir id hardcode change

* Fix framework parse

---------

Co-authored-by: Jinpeng Qi <[email protected]>

Author: jinpengqi

src/sagemaker/inference_recommender/inference_recommender_mixin.py

@@ -100,13 +100,15 @@ def right_size(
                         'OMP_NUM_THREADS': CategoricalParameter(['1', '2', '3', '4'])
                     }]
 
-            phases (list[Phase]): Specifies the criteria for increasing load
-                during endpoint load tests. (default: None).
-            traffic_type (str): Specifies the traffic type that matches the phases. (default: None).
-            max_invocations (str): defines invocation limit for endpoint load tests (default: None).
-            model_latency_thresholds (list[ModelLatencyThreshold]): defines the response latency
-                thresholds for endpoint load tests (default: None).
-            max_tests (int): restricts how many endpoints are allowed to be
+            phases (list[Phase]): Shape of the traffic pattern to use in the load test
+                (default: None).
+            traffic_type (str): Specifies the traffic pattern type. Currently only supports
+                one type 'PHASES' (default: None).
+            max_invocations (str): defines the minimum invocations per minute for the endpoint
+                to support (default: None).
+            model_latency_thresholds (list[ModelLatencyThreshold]): defines the maximum response
+                latency for endpoints to support (default: None).
+            max_tests (int): restricts how many endpoints in total are allowed to be
                 spun up for this job (default: None).
             max_parallel_tests (int): restricts how many concurrent endpoints
                 this job is allowed to spin up (default: None).
@@ -121,7 +123,7 @@ def right_size(
             raise ValueError("right_size() is currently only supported with a registered model")
 
         if not framework and self._framework():
-            framework = INFERENCE_RECOMMENDER_FRAMEWORK_MAPPING.get(self._framework, framework)
+            framework = INFERENCE_RECOMMENDER_FRAMEWORK_MAPPING.get(self._framework(), framework)
 
         framework_version = self._get_framework_version()
 
@@ -177,15 +179,16 @@ def right_size(
 
     def _update_params(
         self,
-        instance_type,
-        initial_instance_count,
-        accelerator_type,
-        async_inference_config,
-        serverless_inference_config,
-        inference_recommendation_id,
-        inference_recommender_job_results,
+        **kwargs,
     ):
         """Check and update params based on inference recommendation id or right size case"""
+        instance_type = kwargs["instance_type"]
+        initial_instance_count = kwargs["initial_instance_count"]
+        accelerator_type = kwargs["accelerator_type"]
+        async_inference_config = kwargs["async_inference_config"]
+        serverless_inference_config = kwargs["serverless_inference_config"]
+        inference_recommendation_id = kwargs["inference_recommendation_id"]
+        inference_recommender_job_results = kwargs["inference_recommender_job_results"]
         if inference_recommendation_id is not None:
             inference_recommendation = self._update_params_for_recommendation_id(
                 instance_type=instance_type,
@@ -394,7 +397,6 @@ def _update_params_for_recommendation_id(
                 self.model_data = model_res["PrimaryContainer"]["ModelDataUrl"]
                 self.image_uri = model_res["PrimaryContainer"]["Image"]
         else:
-            # Update params with compilation recommendation results
             if "InferenceSpecificationName" in model_config:
                 modelpkg_res = sage_client.describe_model_package(
                     ModelPackageName=input_config["ModelPackageVersionArn"]

src/sagemaker/model.py

@@ -1129,6 +1129,15 @@ def deploy(
         """
         removed_kwargs("update_endpoint", kwargs)
 
+        self._init_sagemaker_session_if_does_not_exist(instance_type)
+
+        tags = add_jumpstart_tags(
+            tags=tags, inference_model_uri=self.model_data, inference_script_uri=self.source_dir
+        )
+
+        if self.role is None:
+            raise ValueError("Role can not be null for deploying a model")
+
         if (
             inference_recommendation_id is not None
             or self.inference_recommender_job_results is not None
@@ -1143,15 +1152,6 @@ def deploy(
                 inference_recommender_job_results=self.inference_recommender_job_results,
             )
 
-        self._init_sagemaker_session_if_does_not_exist(instance_type)
-
-        tags = add_jumpstart_tags(
-            tags=tags, inference_model_uri=self.model_data, inference_script_uri=self.source_dir
-        )
-
-        if self.role is None:
-            raise ValueError("Role can not be null for deploying a model")
-
         is_async = async_inference_config is not None
         if is_async and not isinstance(async_inference_config, AsyncInferenceConfig):
             raise ValueError("async_inference_config needs to be a AsyncInferenceConfig object")

tests/integ/test_inference_recommender.py

@@ -43,6 +43,7 @@ def default_right_sized_model(sagemaker_session, cpu_instance_type):
     with timeout(minutes=45):
         try:
             model_package_group_name = unique_name_from_base("test-ir-right-size-model-pkg-sklearn")
+            ir_job_name = unique_name_from_base("test-ir-right-size-job-name")
             model_data = sagemaker_session.upload_data(path=IR_SKLEARN_MODEL)
             payload_data = sagemaker_session.upload_data(path=IR_SKLEARN_PAYLOAD)
 
@@ -66,13 +67,15 @@ def default_right_sized_model(sagemaker_session, cpu_instance_type):
 
             return (
                 sklearn_model_package.right_size(
+                    job_name=ir_job_name,
                     sample_payload_url=payload_data,
                     supported_content_types=IR_SKLEARN_CONTENT_TYPE,
                     supported_instance_types=[cpu_instance_type],
                     framework=IR_SKLEARN_FRAMEWORK,
                     log_level="Quiet",
                 ),
                 model_package_group_name,
+                ir_job_name,
             )
         except Exception:
             sagemaker_session.sagemaker_client.delete_model_package(
@@ -157,7 +160,7 @@ def test_default_right_size_and_deploy_registered_model_sklearn(
 ):
     endpoint_name = unique_name_from_base("test-ir-right-size-default-sklearn")
 
-    right_size_model_package, model_package_group_name = default_right_sized_model
+    right_size_model_package, model_package_group_name, ir_job_name = default_right_sized_model
     with timeout(minutes=45):
         try:
             right_size_model_package.predictor_cls = SKLearnPredictor
@@ -169,12 +172,6 @@ def test_default_right_size_and_deploy_registered_model_sklearn(
             assert inference is not None
             assert 26 == len(inference)
         finally:
-            sagemaker_session.sagemaker_client.delete_model_package(
-                ModelPackageName=right_size_model_package.model_package_arn
-            )
-            sagemaker_session.sagemaker_client.delete_model_package_group(
-                ModelPackageGroupName=model_package_group_name
-            )
             predictor.delete_model()
             predictor.delete_endpoint()
 
@@ -209,3 +206,35 @@ def test_advanced_right_size_and_deploy_registered_model_sklearn(
 
 # TODO when we've added support for inference_recommendation_id
 # then add tests to test Framework models
+@pytest.mark.slow_test
+def test_deploy_inference_recommendation_id_with_registered_model_sklearn(
+    default_right_sized_model, sagemaker_session
+):
+    right_size_model_package, model_package_group_name, ir_job_name = default_right_sized_model
+    endpoint_name = unique_name_from_base("test-rec-id-deployment-default-sklearn")
+    rec_res = sagemaker_session.sagemaker_client.describe_inference_recommendations_job(
+        JobName=ir_job_name
+    )
+    rec_id = rec_res["InferenceRecommendations"][0]["RecommendationId"]
+
+    with timeout(minutes=45):
+        try:
+            right_size_model_package.predictor_cls = SKLearnPredictor
+            predictor = right_size_model_package.deploy(
+                inference_recommendation_id=rec_id, endpoint_name=endpoint_name
+            )
+
+            payload = pd.read_csv(IR_SKLEARN_DATA, header=None)
+
+            inference = predictor.predict(payload)
+            assert inference is not None
+            assert 26 == len(inference)
+        finally:
+            sagemaker_session.sagemaker_client.delete_model_package(
+                ModelPackageName=right_size_model_package.model_package_arn
+            )
+            sagemaker_session.sagemaker_client.delete_model_package_group(
+                ModelPackageGroupName=model_package_group_name
+            )
+            predictor.delete_model()
+            predictor.delete_endpoint()

tests/unit/sagemaker/model/test_deploy.py

@@ -530,12 +530,17 @@ def test_deploy_with_wrong_recommendation_id(sagemaker_session):
         )
 
 
+def mock_describe_model_package(ModelPackageName):
+    if ModelPackageName == IR_MODEL_PACKAGE_VERSION_ARN:
+        return DESCRIBE_MODEL_PACKAGE_RESPONSE
+
+
 def test_deploy_with_recommendation_id_with_model_pkg_arn(sagemaker_session):
     sagemaker_session.sagemaker_client.describe_inference_recommendations_job.return_value = (
         create_inference_recommendations_job_default_with_model_package_arn()
     )
-    sagemaker_session.sagemaker_client.describe_model_package.return_value = (
-        DESCRIBE_MODEL_PACKAGE_RESPONSE
+    sagemaker_session.sagemaker_client.describe_model_package.side_effect = (
+        mock_describe_model_package
     )
 
     model = Model(MODEL_IMAGE, MODEL_DATA, sagemaker_session=sagemaker_session, role=ROLE)
@@ -544,29 +549,27 @@ def test_deploy_with_recommendation_id_with_model_pkg_arn(sagemaker_session):
         inference_recommendation_id=RECOMMENDATION_ID,
     )
 
-    sagemaker_session.sagemaker_client.describe_model_package.assert_called_once_with(
-        ModelPackageName=IR_MODEL_PACKAGE_VERSION_ARN
-    )
     assert model.model_data == IR_MODEL_DATA
     assert model.image_uri == IR_IMAGE
     assert model.env == IR_ENV
 
 
 def test_deploy_with_recommendation_id_with_model_name(sagemaker_session):
+    def mock_describe_model(ModelName):
+        if ModelName == IR_MODEL_NAME:
+            return DESCRIBE_MODEL_RESPONSE
+
     sagemaker_session.sagemaker_client.describe_inference_recommendations_job.return_value = (
         create_inference_recommendations_job_default_with_model_name()
     )
-    sagemaker_session.sagemaker_client.describe_model.return_value = DESCRIBE_MODEL_RESPONSE
+    sagemaker_session.sagemaker_client.describe_model.side_effect = mock_describe_model
 
     model = Model(MODEL_IMAGE, MODEL_DATA, sagemaker_session=sagemaker_session, role=ROLE)
 
     model.deploy(
         inference_recommendation_id=RECOMMENDATION_ID,
     )
 
-    sagemaker_session.sagemaker_client.describe_model.assert_called_once_with(
-        ModelName=IR_MODEL_NAME
-    )
     assert model.model_data == IR_MODEL_DATA
     assert model.image_uri == IR_IMAGE
     assert model.env == IR_ENV
@@ -576,8 +579,8 @@ def test_deploy_with_recommendation_id_with_model_pkg_arn_and_compilation(sagema
     sagemaker_session.sagemaker_client.describe_inference_recommendations_job.return_value = (
         create_inference_recommendations_job_default_with_model_package_arn_and_compilation()
     )
-    sagemaker_session.sagemaker_client.describe_model_package.return_value = (
-        DESCRIBE_MODEL_PACKAGE_RESPONSE
+    sagemaker_session.sagemaker_client.describe_model_package.side_effect = (
+        mock_describe_model_package
     )
 
     model = Model(MODEL_IMAGE, MODEL_DATA, sagemaker_session=sagemaker_session, role=ROLE)
@@ -586,19 +589,20 @@ def test_deploy_with_recommendation_id_with_model_pkg_arn_and_compilation(sagema
         inference_recommendation_id=RECOMMENDATION_ID,
     )
 
-    sagemaker_session.sagemaker_client.describe_model_package.assert_called_once_with(
-        ModelPackageName=IR_MODEL_PACKAGE_VERSION_ARN
-    )
     assert model.model_data == IR_COMPILATION_MODEL_DATA
     assert model.image_uri == IR_COMPILATION_IMAGE
 
 
 def test_deploy_with_recommendation_id_with_model_name_and_compilation(sagemaker_session):
+    def mock_describe_compilation_job(CompilationJobName):
+        if CompilationJobName == IR_COMPILATION_JOB_NAME:
+            return DESCRIBE_COMPILATION_JOB_RESPONSE
+
     sagemaker_session.sagemaker_client.describe_inference_recommendations_job.return_value = (
         create_inference_recommendations_job_default_with_model_name_and_compilation()
     )
-    sagemaker_session.sagemaker_client.describe_compilation_job.return_value = (
-        DESCRIBE_COMPILATION_JOB_RESPONSE
+    sagemaker_session.sagemaker_client.describe_compilation_job.side_effect = (
+        mock_describe_compilation_job
     )
 
     model = Model(MODEL_IMAGE, MODEL_DATA, sagemaker_session=sagemaker_session, role=ROLE)
@@ -607,9 +611,6 @@ def test_deploy_with_recommendation_id_with_model_name_and_compilation(sagemaker
         inference_recommendation_id=RECOMMENDATION_ID,
     )
 
-    sagemaker_session.sagemaker_client.describe_compilation_job.assert_called_once_with(
-        CompilationJobName=IR_COMPILATION_JOB_NAME
-    )
     assert model.model_data == IR_COMPILATION_MODEL_DATA
     assert model.image_uri == IR_COMPILATION_IMAGE