Add deployment support for instant recommendations

Author: jinpengqi

src/sagemaker/inference_recommender/inference_recommender_mixin.py

@@ -365,12 +365,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:
@@ -390,26 +384,66 @@ def _update_params_for_recommendation_id(
         recommendation_job_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
-        )
-        input_config = recommendation_res["InputConfig"]
 
-        recommendation = next(
-            (
-                rec
-                for rec in recommendation_res["InferenceRecommendations"]
-                if rec["RecommendationId"] == inference_recommendation_id
-            ),
-            None,
-        )
+        # Retrieve model or inference recommendation job details
+        recommendation_res, model_res = None, None
+        try:
+            recommendation_res = sage_client.describe_inference_recommendations_job(
+                JobName=recommendation_job_name
+            )
+        except sage_client.exceptions.ResourceNotFound:
+            pass
+        try:
+            model_res = sage_client.describe_model(ModelName=recommendation_job_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")
 
-        if not recommendation:
-            raise ValueError(
-                "inference_recommendation_id does not exist in InferenceRecommendations list"
+        # Search the recommendation from above describe result lists
+        inference_recommendation, instant_recommendation = None, None
+        if recommendation_res:
+            inference_recommendation = next(
+                (
+                    rec
+                    for rec in recommendation_res["InferenceRecommendations"]
+                    if rec["RecommendationId"] == inference_recommendation_id
+                ),
+                None,
             )
+        if model_res:
+            instant_recommendation = next(
+                (
+                    rec
+                    for rec in model_res["DeploymentRecommendation"][
+                        "RealTimeInferenceRecommendations"
+                    ]
+                    if rec["RecommendationId"] == inference_recommendation_id
+                ),
+                None,
+            )
+        if inference_recommendation is None and instant_recommendation is None:
+            raise ValueError("Inference Recommendation id does not exist")
+
+        # Update params beased on instant recommendation
+        if instant_recommendation:
+            if initial_instance_count is None:
+                raise ValueError(
+                    "Please specify initial_instance_count with instant recommendation id"
+                )
+            self.env.update(instant_recommendation["Environment"])
+            instance_type = instant_recommendation["InstanceType"]
+            return (instance_type, initial_instance_count)
 
-        model_config = recommendation["ModelConfiguration"]
+        # Update params based on default inference recommendation
+        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."
+            )
+        input_config = recommendation_res["InputConfig"]
+        model_config = inference_recommendation["ModelConfiguration"]
         envs = (
             model_config["EnvironmentParameters"]
             if "EnvironmentParameters" in model_config
@@ -458,8 +492,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 = inference_recommendation["EndpointConfiguration"]["InstanceType"]
+        initial_instance_count = inference_recommendation["EndpointConfiguration"][
+            "InitialInstanceCount"
+        ]
 
         return (instance_type, initial_instance_count)
 

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_INSTANT_RECOMMENDATION_ID = IR_MODEL_NAME + "/ad3ec9ee"
 RECOMMENDATION_ID = IR_JOB_NAME + "/5bcee92e"
+INSTANT_RECOMMENDATION_ID = IR_MODEL_NAME + "/v0KObO5d"
+INSTANT_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": INSTANT_RECOMMENDATION_ID,
+                "InstanceType": "ml.g4dn.2xlarge",
+                "Environment": INSTANT_RECOMMENDATION_ENV,
+            },
+            {
+                "RecommendationId": "test-model-name/d248qVYU",
+                "InstanceType": "ml.c6i.large",
+                "Environment": {},
+            },
+        ],
+    },
 }
 
 DESCRIBE_MODEL_PACKAGE_RESPONSE = {

tests/unit/sagemaker/model/test_deploy.py

@@ -26,6 +26,8 @@
     DESCRIBE_COMPILATION_JOB_RESPONSE,
     DESCRIBE_MODEL_PACKAGE_RESPONSE,
     DESCRIBE_MODEL_RESPONSE,
+    INSTANT_RECOMMENDATION_ENV,
+    INSTANT_RECOMMENDATION_ID,
     INVALID_RECOMMENDATION_ID,
     IR_COMPILATION_JOB_NAME,
     IR_ENV,
@@ -35,6 +37,7 @@
     IR_MODEL_PACKAGE_VERSION_ARN,
     IR_COMPILATION_IMAGE,
     IR_COMPILATION_MODEL_DATA,
+    NOT_EXISTED_INSTANT_RECOMMENDATION_ID,
     RECOMMENDATION_ID,
     NOT_EXISTED_RECOMMENDATION_ID,
 )
@@ -543,6 +546,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(
@@ -553,7 +561,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(
@@ -564,7 +572,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(
@@ -615,6 +623,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)
 
@@ -627,11 +636,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()
     )
@@ -655,6 +665,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)
 
@@ -677,6 +688,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)
 
@@ -695,18 +707,49 @@ def test_deploy_with_not_existed_recommendation_id(sagemaker_session):
     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 does not exist",
     ):
         model.deploy(
             inference_recommendation_id=NOT_EXISTED_RECOMMENDATION_ID,
         )
 
 
+def test_deploy_with_invalid_instant_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 does not exist",
+    ):
+        model.deploy(
+            inference_recommendation_id=NOT_EXISTED_INSTANT_RECOMMENDATION_ID,
+        )
+
+
+def test_deploy_with_valid_instant_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=INSTANT_RECOMMENDATION_ID,
+        initial_instance_count=INSTANCE_COUNT,
+    )
+
+    assert model.model_data == MODEL_DATA
+    assert model.image_uri == MODEL_IMAGE
+    assert model.env == INSTANT_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())