fix: pass in inference_ami_version to model_based endpoint type

Author: timkuo-amazon

src/sagemaker/model.py

@@ -1492,6 +1492,9 @@ def deploy(
                     }
             model_reference_arn (Optional [str]): Hub Content Arn of a Model Reference type
                 content (default: None).
+            inference_ami_version (Optional [str]): Specifies an option from a collection of preconfigured
+             Amazon Machine Image (AMI) images. For a full list of options, see:
+             https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_ProductionVariant.html
         Raises:
              ValueError: If arguments combination check failed in these circumstances:
                 - If no role is specified or
@@ -1743,6 +1746,7 @@ def deploy(
                 model_data_download_timeout=model_data_download_timeout,
                 container_startup_health_check_timeout=container_startup_health_check_timeout,
                 routing_config=routing_config,
+                inference_ami_version=inference_ami_version,
             )
             if endpoint_name:
                 self.endpoint_name = endpoint_name

tests/unit/sagemaker/model/test_deploy.py

@@ -130,6 +130,7 @@ def test_deploy(name_from_base, prepare_container_def, production_variant, sagem
         model_data_download_timeout=None,
         container_startup_health_check_timeout=None,
         routing_config=None,
+        inference_ami_version=None,
     )
 
     sagemaker_session.create_model.assert_called_with(
@@ -192,6 +193,7 @@ def test_deploy_accelerator_type(
         model_data_download_timeout=None,
         container_startup_health_check_timeout=None,
         routing_config=None,
+        inference_ami_version=None,
     )
 
     sagemaker_session.endpoint_from_production_variants.assert_called_with(
@@ -519,6 +521,7 @@ def test_deploy_serverless_inference(production_variant, create_sagemaker_model,
         model_data_download_timeout=None,
         container_startup_health_check_timeout=None,
         routing_config=None,
+        inference_ami_version=None,
     )
 
     sagemaker_session.endpoint_from_production_variants.assert_called_with(
@@ -956,6 +959,7 @@ def test_deploy_customized_volume_size_and_timeout(
         model_data_download_timeout=model_data_download_timeout_sec,
         container_startup_health_check_timeout=startup_health_check_timeout_sec,
         routing_config=None,
+        inference_ami_version=None,
     )
 
     sagemaker_session.create_model.assert_called_with(
@@ -1006,6 +1010,7 @@ def test_deploy_with_resources(sagemaker_session, name_from_base, production_var
         model_data_download_timeout=None,
         container_startup_health_check_timeout=None,
         routing_config=None,
+        inference_ami_version=None,
     )
     sagemaker_session.endpoint_from_production_variants.assert_called_with(
         name=name_from_base(MODEL_NAME),