fix: fix flaky Inference Recommender integration tests (aws#4156)

Co-authored-by: Gary Wang <[email protected]>
Co-authored-by: Joseph Zhang <[email protected]>

Author: 3 people

src/sagemaker/inference_recommender/inference_recommender_mixin.py

@@ -302,11 +302,7 @@ def _update_params_for_right_size(
             )
             return None
 
-        instance_type = self.inference_recommendations[0]["EndpointConfiguration"]["InstanceType"]
-        initial_instance_count = self.inference_recommendations[0]["EndpointConfiguration"][
-            "InitialInstanceCount"
-        ]
-        return (instance_type, initial_instance_count)
+        return self._filter_recommendations_for_realtime()
 
     def _update_params_for_recommendation_id(
         self,
@@ -610,3 +606,16 @@ def _search_recommendation(self, recommendation_list, inference_recommendation_i
             ),
             None,
         )
+
+    # TODO: until we have bandwidth to integrate right_size + deploy with serverless
+    def _filter_recommendations_for_realtime(self):
+        """Filter recommendations list to find a realtime instance"""
+        instance_type = None
+        initial_instance_count = None
+        for recommendations in self.inference_recommendations:
+            if "ServerlessConfig" not in recommendations["EndpointConfiguration"]:
+                instance_type = recommendations["EndpointConfiguration"]["InstanceType"]
+                initial_instance_count = recommendations["EndpointConfiguration"][
+                    "InitialInstanceCount"
+                ]
+        return (instance_type, initial_instance_count)

tests/integ/test_inference_recommender.py

@@ -460,7 +460,8 @@ def test_deploy_inference_recommendation_id_with_registered_model_sklearn(
     rec_res = sagemaker_session.sagemaker_client.describe_inference_recommendations_job(
         JobName=ir_job_name
     )
-    rec_id = rec_res["InferenceRecommendations"][0]["RecommendationId"]
+
+    rec_id = get_realtime_recommendation_id(recommendation_list=rec_res["InferenceRecommendations"])
 
     with timeout(minutes=45):
         try:
@@ -537,3 +538,11 @@ def poll_for_deployment_recommendation(created_base_model, sagemaker_session):
         except Exception as e:
             created_base_model.delete_model()
             raise e
+
+
+def get_realtime_recommendation_id(recommendation_list):
+    """Search recommendation based on recommendation id"""
+    next(
+        (rec["RecommendationId"] for rec in recommendation_list if "InstanceType" in rec),
+        None,
+    )