feat: Add deployment support for deployment recommendations

src/sagemaker/inference_recommender/inference_recommender_mixin.py

@@ -351,12 +351,6 @@ def _update_params_for_recommendation_id(
             return (instance_type, initial_instance_count)
 
         # Validate non-compatible parameters with recommendation id
-        if bool(instance_type) != bool(initial_instance_count):
-            raise ValueError(
-                "Please either do not specify instance_type and initial_instance_count"
-                "since they are in recommendation, or specify both of them if you want"
-                "to override the recommendation."
-            )
         if accelerator_type is not None:
             raise ValueError("accelerator_type is not compatible with inference_recommendation_id.")
         if async_inference_config is not None:
@@ -371,29 +365,45 @@ def _update_params_for_recommendation_id(
         # Validate recommendation id
         if not re.match(r"[a-zA-Z0-9](-*[a-zA-Z0-9]){0,63}\/\w{8}$", inference_recommendation_id):
             raise ValueError("Inference Recommendation id is not valid")
-        recommendation_job_name = inference_recommendation_id.split("/")[0]
+        job_or_model_name = inference_recommendation_id.split("/")[0]
 
         sage_client = self.sagemaker_session.sagemaker_client
-        recommendation_res = sage_client.describe_inference_recommendations_job(
-            JobName=recommendation_job_name
+
+        # Desribe inference recommendation job and model details
+        recommendation_res, model_res = self._describe_recommendation_job_and_model(
+            sage_client=sage_client,
+            job_or_model_name=job_or_model_name,
         )
-        input_config = recommendation_res["InputConfig"]
 
-        recommendation = next(
-            (
-                rec
-                for rec in recommendation_res["InferenceRecommendations"]
-                if rec["RecommendationId"] == inference_recommendation_id
-            ),
-            None,
+        # Search the recommendation from above describe results
+        (
+            right_size_recommendation,
+            model_recommendation,
+        ) = self._get_right_size_and_model_recommendation(
+            recommendation_res=recommendation_res,
+            model_res=model_res,
+            inference_recommendation_id=inference_recommendation_id,
         )
 
-        if not recommendation:
+        # Update params beased on model recommendation
+        if model_recommendation:
+            if initial_instance_count is None:
+                raise ValueError(
+                    "Please specify initial_instance_count with model recommendation id"
+                )
+            self.env.update(model_recommendation["Environment"])
+            instance_type = model_recommendation["InstanceType"]
+            return (instance_type, initial_instance_count)
+
+        # Update params based on default inference recommendation
+        if bool(instance_type) != bool(initial_instance_count):
             raise ValueError(
-                "inference_recommendation_id does not exist in InferenceRecommendations list"
+                "Please either do not specify instance_type and initial_instance_count"
+                "since they are in recommendation, or specify both of them if you want"
+                "to override the recommendation."
             )
-
-        model_config = recommendation["ModelConfiguration"]
+        input_config = recommendation_res["InputConfig"]
+        model_config = right_size_recommendation["ModelConfiguration"]
         envs = (
             model_config["EnvironmentParameters"]
             if "EnvironmentParameters" in model_config
@@ -442,8 +452,10 @@ def _update_params_for_recommendation_id(
                 self.model_data = compilation_res["ModelArtifacts"]["S3ModelArtifacts"]
                 self.image_uri = compilation_res["InferenceImage"]
 
-        instance_type = recommendation["EndpointConfiguration"]["InstanceType"]
-        initial_instance_count = recommendation["EndpointConfiguration"]["InitialInstanceCount"]
+        instance_type = right_size_recommendation["EndpointConfiguration"]["InstanceType"]
+        initial_instance_count = right_size_recommendation["EndpointConfiguration"][
+            "InitialInstanceCount"
+        ]
 
         return (instance_type, initial_instance_count)
 
@@ -513,3 +525,57 @@ def _convert_to_stopping_conditions_json(
                 threshold.to_json for threshold in model_latency_thresholds
             ]
         return stopping_conditions
+
+    def _get_right_size_and_model_recommendation(
+        self,
+        model_res=None,
+        recommendation_res=None,
+        inference_recommendation_id=None,
+    ):
+        """Get recommendation from right size job or model"""
+        right_size_recommendation, model_recommendation = None, None
+        if recommendation_res:
+            right_size_recommendation = self._get_recommendation(
+                recommendation_list=recommendation_res["InferenceRecommendations"],
+                inference_recommendation_id=inference_recommendation_id,
+            )
+        if model_res:
+            model_recommendation = self._get_recommendation(
+                recommendation_list=model_res["DeploymentRecommendation"][
+                    "RealTimeInferenceRecommendations"
+                ],
+                inference_recommendation_id=inference_recommendation_id,
+            )
+        if right_size_recommendation is None and model_recommendation is None:
+            raise ValueError("Inference Recommendation id is not valid")
+
+        return right_size_recommendation, model_recommendation
+
+    def _get_recommendation(self, recommendation_list, inference_recommendation_id):
+        """Get recommendation based on recommendation id"""
+        return next(
+            (
+                rec
+                for rec in recommendation_list
+                if rec["RecommendationId"] == inference_recommendation_id
+            ),
+            None,
+        )
+
+    def _describe_recommendation_job_and_model(self, sage_client, job_or_model_name):
+        """Describe inference recommendation job and model results"""
+        recommendation_res, model_res = None, None
+        try:
+            recommendation_res = sage_client.describe_inference_recommendations_job(
+                JobName=job_or_model_name
+            )
+        except sage_client.exceptions.ResourceNotFound:
+            pass
+        try:
+            model_res = sage_client.describe_model(ModelName=job_or_model_name)
+        except sage_client.exceptions.ResourceNotFound:
+            pass
+        if recommendation_res is None and model_res is None:
+            raise ValueError("Inference Recommendation id is not valid")
+
+        return recommendation_res, model_res

tests/unit/sagemaker/inference_recommender/constants.py

@@ -37,7 +37,10 @@
 
 INVALID_RECOMMENDATION_ID = "ir-job6ab0ff22"
 NOT_EXISTED_RECOMMENDATION_ID = IR_JOB_NAME + "/ad3ec9ee"
+NOT_EXISTED_MODEL_RECOMMENDATION_ID = IR_MODEL_NAME + "/ad3ec9ee"
 RECOMMENDATION_ID = IR_JOB_NAME + "/5bcee92e"
+MODEL_RECOMMENDATION_ID = IR_MODEL_NAME + "/v0KObO5d"
+MODEL_RECOMMENDATION_ENV = {"TS_DEFAULT_WORKERS_PER_MODEL": "4"}
 
 IR_CONTAINER_CONFIG = {
     "Domain": "MACHINE_LEARNING",
@@ -95,6 +98,21 @@
         "Image": IR_IMAGE,
         "ModelDataUrl": IR_MODEL_DATA,
     },
+    "DeploymentRecommendation": {
+        "RecommendationStatus": "COMPLETED",
+        "RealTimeInferenceRecommendations": [
+            {
+                "RecommendationId": MODEL_RECOMMENDATION_ID,
+                "InstanceType": "ml.g4dn.2xlarge",
+                "Environment": MODEL_RECOMMENDATION_ENV,
+            },
+            {
+                "RecommendationId": "test-model-name/d248qVYU",
+                "InstanceType": "ml.c6i.large",
+                "Environment": {},
+            },
+        ],
+    },
 }
 
 DESCRIBE_MODEL_PACKAGE_RESPONSE = {

tests/unit/sagemaker/model/test_deploy.py

@@ -25,6 +25,8 @@
     DESCRIBE_COMPILATION_JOB_RESPONSE,
     DESCRIBE_MODEL_PACKAGE_RESPONSE,
     DESCRIBE_MODEL_RESPONSE,
+    MODEL_RECOMMENDATION_ENV,
+    MODEL_RECOMMENDATION_ID,
     INVALID_RECOMMENDATION_ID,
     IR_COMPILATION_JOB_NAME,
     IR_ENV,
@@ -34,6 +36,7 @@
     IR_MODEL_PACKAGE_VERSION_ARN,
     IR_COMPILATION_IMAGE,
     IR_COMPILATION_MODEL_DATA,
+    NOT_EXISTED_MODEL_RECOMMENDATION_ID,
     RECOMMENDATION_ID,
     NOT_EXISTED_RECOMMENDATION_ID,
 )
@@ -470,6 +473,11 @@ def test_deploy_wrong_async_inferenc_config(sagemaker_session):
 
 
 def test_deploy_ir_with_incompatible_parameters(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.return_value = None
+
     model = Model(MODEL_IMAGE, MODEL_DATA, sagemaker_session=sagemaker_session, role=ROLE)
 
     with pytest.raises(
@@ -480,7 +488,7 @@ def test_deploy_ir_with_incompatible_parameters(sagemaker_session):
     ):
         model.deploy(
             instance_type=INSTANCE_TYPE,
-            inference_recommendation_id=INFERENCE_RECOMMENDATION_ID,
+            inference_recommendation_id=RECOMMENDATION_ID,
         )
 
     with pytest.raises(
@@ -491,7 +499,7 @@ def test_deploy_ir_with_incompatible_parameters(sagemaker_session):
     ):
         model.deploy(
             initial_instance_count=INSTANCE_COUNT,
-            inference_recommendation_id=INFERENCE_RECOMMENDATION_ID,
+            inference_recommendation_id=RECOMMENDATION_ID,
         )
 
     with pytest.raises(
@@ -542,6 +550,7 @@ def test_deploy_with_recommendation_id_with_model_pkg_arn(sagemaker_session):
     sagemaker_session.sagemaker_client.describe_model_package.side_effect = (
         mock_describe_model_package
     )
+    sagemaker_session.sagemaker_client.describe_model.return_value = None
 
     model = Model(MODEL_IMAGE, MODEL_DATA, sagemaker_session=sagemaker_session, role=ROLE)
 
@@ -554,11 +563,12 @@ def test_deploy_with_recommendation_id_with_model_pkg_arn(sagemaker_session):
     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
+def mock_describe_model(ModelName):
+    if ModelName == IR_MODEL_NAME:
+        return DESCRIBE_MODEL_RESPONSE
+
 
+def test_deploy_with_recommendation_id_with_model_name(sagemaker_session):
     sagemaker_session.sagemaker_client.describe_inference_recommendations_job.return_value = (
         create_inference_recommendations_job_default_with_model_name()
     )
@@ -582,6 +592,7 @@ def test_deploy_with_recommendation_id_with_model_pkg_arn_and_compilation(sagema
     sagemaker_session.sagemaker_client.describe_model_package.side_effect = (
         mock_describe_model_package
     )
+    sagemaker_session.sagemaker_client.describe_model.return_value = None
 
     model = Model(MODEL_IMAGE, MODEL_DATA, sagemaker_session=sagemaker_session, role=ROLE)
 
@@ -604,6 +615,7 @@ def mock_describe_compilation_job(CompilationJobName):
     sagemaker_session.sagemaker_client.describe_compilation_job.side_effect = (
         mock_describe_compilation_job
     )
+    sagemaker_session.sagemaker_client.describe_model.side_effect = mock_describe_model
 
     model = Model(MODEL_IMAGE, MODEL_DATA, sagemaker_session=sagemaker_session, role=ROLE)
 
@@ -615,25 +627,56 @@ def mock_describe_compilation_job(CompilationJobName):
     assert model.image_uri == IR_COMPILATION_IMAGE
 
 
-def test_deploy_with_not_existed_recommendation_id(sagemaker_session):
+def test_deploy_with_invalid_inference_recommendation_id(sagemaker_session):
     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_model.return_value = None
 
     model = Model(MODEL_IMAGE, MODEL_DATA, sagemaker_session=sagemaker_session, role=ROLE)
 
     with pytest.raises(
         ValueError,
-        match="inference_recommendation_id does not exist in InferenceRecommendations list",
+        match="Inference Recommendation id is not valid",
     ):
         model.deploy(
             inference_recommendation_id=NOT_EXISTED_RECOMMENDATION_ID,
         )
 
 
+def test_deploy_with_invalid_model_recommendation_id(sagemaker_session):
+    sagemaker_session.sagemaker_client.describe_inference_recommendations_job.return_value = None
+    sagemaker_session.sagemaker_client.describe_model.side_effect = mock_describe_model
+
+    model = Model(MODEL_IMAGE, MODEL_DATA, sagemaker_session=sagemaker_session, role=ROLE)
+
+    with pytest.raises(
+        ValueError,
+        match="Inference Recommendation id is not valid",
+    ):
+        model.deploy(
+            inference_recommendation_id=NOT_EXISTED_MODEL_RECOMMENDATION_ID,
+        )
+
+
+def test_deploy_with_valid_model_recommendation_id(sagemaker_session):
+    sagemaker_session.sagemaker_client.describe_inference_recommendations_job.return_value = None
+    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=MODEL_RECOMMENDATION_ID,
+        initial_instance_count=INSTANCE_COUNT,
+    )
+
+    assert model.model_data == MODEL_DATA
+    assert model.image_uri == MODEL_IMAGE
+    assert model.env == MODEL_RECOMMENDATION_ENV
+
+
 @patch("sagemaker.model.Model._create_sagemaker_model", Mock())
 @patch("sagemaker.predictor.Predictor._get_endpoint_config_name", Mock())
 @patch("sagemaker.predictor.Predictor._get_model_names", Mock())